From ee7c315afb7eb55e97ffaa061355ca29c9fd9406 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 14:00:16 -0400 Subject: [PATCH 01/44] feat(datasets): scaffold nemo-datasets plugin Add a new nemo-datasets plugin (non-member scaffold, like example-plugin) that registers the `nemo datasets` CLI via the nemo.cli entry point. The profile command surface is wired with a stub; the profiling core lands in follow-up commits. Signed-off-by: Albert Cui --- plugins/nemo-datasets/pyproject.toml | 29 ++++++++++++++++++ .../src/nemo_datasets_plugin/__init__.py | 8 +++++ .../src/nemo_datasets_plugin/cli.py | 30 +++++++++++++++++++ plugins/nemo-datasets/tests/test_cli.py | 22 ++++++++++++++ pyproject.toml | 3 ++ pytest.ini | 1 + 6 files changed, 93 insertions(+) create mode 100644 plugins/nemo-datasets/pyproject.toml create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py create mode 100644 plugins/nemo-datasets/tests/test_cli.py diff --git a/plugins/nemo-datasets/pyproject.toml b/plugins/nemo-datasets/pyproject.toml new file mode 100644 index 0000000000..89075e4bc0 --- /dev/null +++ b/plugins/nemo-datasets/pyproject.toml @@ -0,0 +1,29 @@ +[project] +name = "nemo-datasets-plugin" +description = "Dataset profiler for NeMo Platform filesets." +requires-python = ">=3.11,<3.15" +dependencies = [ + "nemo-platform-plugin", + "pyarrow>=17.0.0", + "pyyaml>=6.0.2", + "pydantic>=2.10.3", + "typer>=0.20.0,<0.26", +] +version = "0.1.0" + +[tool.uv.sources] +nemo-platform-plugin = { path = "../../packages/nemo_platform_plugin", editable = true } + +[project.entry-points."nemo.cli"] +datasets = "nemo_datasets_plugin.cli:DatasetsCLI" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/nemo_datasets_plugin"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py new file mode 100644 index 0000000000..ecbf54d3aa --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dataset profiler for NeMo Platform filesets. + +Computes a ``DatasetProfile`` (the contract in ``nemo_platform_plugin.files.dataset_profile``) from +a directory of dataset files, and exposes it as the ``nemo datasets`` CLI. +""" diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py new file mode 100644 index 0000000000..0dcbc27343 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The ``nemo datasets`` CLI — registered under the ``nemo.cli`` entry point.""" + +from __future__ import annotations + +import typer +from nemo_platform_plugin.cli import NemoCLI + + +class DatasetsCLI(NemoCLI): + """Exposes dataset commands as ``nemo datasets ...``.""" + + name = "datasets" + description = "Profile datasets stored as filesets." + + def get_cli(self) -> typer.Typer: + app = typer.Typer(help="Dataset profiling commands.") + + @app.command() + def profile( + path: str = typer.Argument(..., help="Path to a local directory of dataset files."), + output: str = typer.Option("json", "--output", "-o", help="Output format: json | yaml."), + ) -> None: + """Profile a local dataset directory and print its DatasetProfile.""" + # The profiling core lands in a follow-up commit; this wires up the command surface. + raise typer.BadParameter("dataset profiling is not implemented yet") + + return app diff --git a/plugins/nemo-datasets/tests/test_cli.py b/plugins/nemo-datasets/tests/test_cli.py new file mode 100644 index 0000000000..65e2d3886e --- /dev/null +++ b/plugins/nemo-datasets/tests/test_cli.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Smoke tests for the ``nemo datasets`` CLI surface.""" + +from nemo_datasets_plugin.cli import DatasetsCLI +from typer.testing import CliRunner + +runner = CliRunner() + + +def test_cli_metadata(): + cli = DatasetsCLI() + assert cli.name == "datasets" + assert cli.description + + +def test_profile_command_registered(): + app = DatasetsCLI().get_cli() + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0 + assert "profile" in result.stdout diff --git a/pyproject.toml b/pyproject.toml index 8d030337b7..a89e3a8e54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -585,6 +585,9 @@ extra-paths = [ # example-plugin is not a workspace member; add its src so ty can # resolve nmp.example_plugin imports when checking plugin test files. "plugins/example-plugin/src", + # nemo-datasets plugin is not a workspace member yet; add its src so ty can + # resolve nemo_datasets_plugin imports when checking plugin test files. + "plugins/nemo-datasets/src", # nemo-guardrails plugin is not a workspace member; add its src so ty can # resolve nemo_guardrails_plugin imports when checking plugin test files. "plugins/nemo-guardrails/src", diff --git a/pytest.ini b/pytest.ini index 0568982260..6fbd695fc8 100644 --- a/pytest.ini +++ b/pytest.ini @@ -9,6 +9,7 @@ python_functions = test_* pythonpath = . plugins/example-plugin/src + plugins/nemo-datasets/src plugins/nemo-deployments/src plugins/nemo-deployments/tests/unit plugins/nemo-deployments/tests/integration From 981bf9ddd7b13677dcc61afab39c76937b1128d5 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 14:00:16 -0400 Subject: [PATCH 02/44] feat(datasets): add file-source seam and parquet/jsonl readers Introduce the profiler core's file-source seam (FileSource protocol + LocalFileSource) and a stateless per-format reader registry with parquet and jsonl readers. Parquet's footer yields an exact row count and declared schema; jsonl reports an exact count only on a full read. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/__init__.py | 13 ++ .../profiler/file_source.py | 56 +++++++++ .../profiler/readers/__init__.py | 28 +++++ .../profiler/readers/base.py | 65 ++++++++++ .../profiler/readers/jsonl.py | 32 +++++ .../profiler/readers/parquet.py | 36 ++++++ plugins/nemo-datasets/tests/test_readers.py | 117 ++++++++++++++++++ 7 files changed, 347 insertions(+) create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py create mode 100644 plugins/nemo-datasets/tests/test_readers.py diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py new file mode 100644 index 0000000000..82bab560ba --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The dataset profiler core. + +A dependency-free library (pyarrow + stdlib) that reads dataset files through a +:class:`~nemo_datasets_plugin.profiler.file_source.FileSource` seam and per-format readers, and — in +later commits — measures and classifies them into a ``DatasetProfile``. +""" + +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource, LocalFileSource + +__all__ = ["FileEntry", "FileSource", "LocalFileSource"] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py new file mode 100644 index 0000000000..4500307e50 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/file_source.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The file-source seam. + +The profiler core reads dataset files only through a :class:`FileSource`, so it never touches a +storage API directly. :class:`LocalFileSource` covers a directory on disk; a ranged-read source over +the Files storage API is a later drop-in behind the same two methods. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, Protocol + + +@dataclass(frozen=True) +class FileEntry: + """A file's identity — everything the profiler needs before it reads the contents.""" + + path: str # POSIX-style path relative to the source root + size_bytes: int + checksum: str | None = None # "sha256:..." when the source reports one; None otherwise + + +class FileSource(Protocol): + """Read-only access to a set of dataset files.""" + + def list_files(self) -> list[FileEntry]: + """Every file in the source, in a stable order.""" + ... + + def open(self, path: str) -> BinaryIO: + """A binary, seekable stream for one file (``path`` as returned by :meth:`list_files`).""" + ... + + +class LocalFileSource: + """A directory of dataset files on the local filesystem.""" + + def __init__(self, root: str | Path) -> None: + self._root = Path(root) + if not self._root.is_dir(): + raise NotADirectoryError(f"{self._root} is not a directory") + + def list_files(self) -> list[FileEntry]: + entries = [ + FileEntry(path=path.relative_to(self._root).as_posix(), size_bytes=path.stat().st_size) + for path in sorted(self._root.rglob("*")) + if path.is_file() + ] + return entries + + def open(self, path: str) -> BinaryIO: + return open(self._root / path, "rb") diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py new file mode 100644 index 0000000000..94ad5b6e52 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Format readers, registered by importing this package. + +Importing the package populates the registry (each reader module self-registers), so +:func:`get_reader` resolves every supported format. +""" + +from nemo_datasets_plugin.profiler.readers.base import ( + FormatReader, + ReadResult, + detect_format, + get_reader, + register_reader, +) + +# Import for side effects: each module calls register_reader() at import time. +from nemo_datasets_plugin.profiler.readers import jsonl as _jsonl # noqa: F401 isort:skip +from nemo_datasets_plugin.profiler.readers import parquet as _parquet # noqa: F401 isort:skip + +__all__ = [ + "FormatReader", + "ReadResult", + "detect_format", + "get_reader", + "register_reader", +] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py new file mode 100644 index 0000000000..ea331aa3ce --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-format reader contract and registry. + +Each reader is a stateless handler for one on-disk format, safe to reuse across files. Readers are +looked up by ``file_format`` (from :func:`detect_format`) so the pipeline never branches on format +itself. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, ClassVar, Protocol + +import pyarrow as pa + +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource + + +@dataclass(frozen=True) +class ReadResult: + """What a format reader returns for one file.""" + + rows: list[dict[str, Any]] # the rows read (a sample, or all of them) + rows_scanned: int # number of rows actually parsed + num_rows: int | None = None # exact total when cheaply known (e.g. a parquet footer), else None + arrow_schema: pa.Schema | None = None # the declared column schema, when the format carries one + + +class FormatReader(Protocol): + """Reads schema and rows for one file format.""" + + file_format: ClassVar[str] + + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: + """Read up to ``row_cap`` rows (all rows when None) plus whatever the format declares cheaply.""" + ... + + +_READERS: dict[str, FormatReader] = {} + + +def register_reader(reader: FormatReader) -> None: + _READERS[reader.file_format] = reader + + +def get_reader(file_format: str) -> FormatReader: + try: + return _READERS[file_format] + except KeyError: + raise KeyError(f"no reader registered for file format {file_format!r}") from None + + +_EXTENSION_FORMATS = { + ".parquet": "parquet", + ".jsonl": "jsonl", + ".ndjson": "jsonl", +} + + +def detect_format(path: str) -> str | None: + """Map a file path to a registered format by extension, or None when unrecognized.""" + return _EXTENSION_FORMATS.get(Path(path).suffix.lower()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py new file mode 100644 index 0000000000..607f231860 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Line-delimited JSON reader. No declared schema; the row count is exact only on a full read.""" + +from __future__ import annotations + +import json + +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource +from nemo_datasets_plugin.profiler.readers.base import ReadResult, register_reader + + +class JsonlReader: + file_format = "jsonl" + + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: + rows: list[dict] = [] + with source.open(entry.path) as stream: + for raw_line in stream: + stripped = raw_line.strip() + if not stripped: # tolerate blank lines between records + continue + rows.append(json.loads(stripped)) + if row_cap is not None and len(rows) >= row_cap: + break + + num_rows = len(rows) if row_cap is None else None + return ReadResult(rows=rows, rows_scanned=len(rows), num_rows=num_rows, arrow_schema=None) + + +register_reader(JsonlReader()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py new file mode 100644 index 0000000000..b360e18e40 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Parquet reader — the footer gives an exact row count and the declared schema for free.""" + +from __future__ import annotations + +import pyarrow.parquet as pq + +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource +from nemo_datasets_plugin.profiler.readers.base import ReadResult, register_reader + + +class ParquetReader: + file_format = "parquet" + + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: + with source.open(entry.path) as stream: + parquet_file = pq.ParquetFile(stream) + num_rows = parquet_file.metadata.num_rows + arrow_schema = parquet_file.schema_arrow + if row_cap == 0: + return ReadResult(rows=[], rows_scanned=0, num_rows=num_rows, arrow_schema=arrow_schema) + + rows: list[dict] = [] + for batch in parquet_file.iter_batches(batch_size=row_cap or 1024): + rows.extend(batch.to_pylist()) + if row_cap is not None and len(rows) >= row_cap: + break + + if row_cap is not None: + rows = rows[:row_cap] + return ReadResult(rows=rows, rows_scanned=len(rows), num_rows=num_rows, arrow_schema=arrow_schema) + + +register_reader(ParquetReader()) diff --git a/plugins/nemo-datasets/tests/test_readers.py b/plugins/nemo-datasets/tests/test_readers.py new file mode 100644 index 0000000000..0c7ffed45b --- /dev/null +++ b/plugins/nemo-datasets/tests/test_readers.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the file-source seam and the per-format readers.""" + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource +from nemo_datasets_plugin.profiler.readers import detect_format, get_reader + +PARQUET_ROWS = [ + {"prompt": "a", "score": 1}, + {"prompt": "b", "score": 2}, + {"prompt": "c", "score": 3}, +] + + +def _write_parquet(path, rows): + pq.write_table(pa.Table.from_pylist(rows), path) + + +# --- file source --------------------------------------------------------------------------------- + + +def test_local_file_source_lists_sorted_with_sizes(tmp_path): + (tmp_path / "b.jsonl").write_text('{"x": 1}\n') + (tmp_path / "a.parquet").write_bytes(b"not-real-parquet") # only listed here, not parsed + sub = tmp_path / "sub" + sub.mkdir() + (sub / "c.jsonl").write_text("{}\n") + + entries = LocalFileSource(tmp_path).list_files() + + assert [e.path for e in entries] == ["a.parquet", "b.jsonl", "sub/c.jsonl"] + assert all(e.size_bytes > 0 for e in entries) + assert all(e.checksum is None for e in entries) # local sources report no checksum + + +def test_local_file_source_open_reads_bytes(tmp_path): + (tmp_path / "f.jsonl").write_text('{"x": 1}\n') + with LocalFileSource(tmp_path).open("f.jsonl") as stream: + assert stream.read() == b'{"x": 1}\n' + + +def test_local_file_source_rejects_non_directory(tmp_path): + target = tmp_path / "f" + target.write_text("x") + with pytest.raises(NotADirectoryError): + LocalFileSource(target) + + +# --- registry ------------------------------------------------------------------------------------ + + +def test_detect_format_by_extension(): + assert detect_format("data/train-00000-of-00003.parquet") == "parquet" + assert detect_format("x.jsonl") == "jsonl" + assert detect_format("x.ndjson") == "jsonl" + assert detect_format("README.md") is None + + +def test_get_reader_unknown_format_raises(): + with pytest.raises(KeyError): + get_reader("arrow") + + +# --- parquet reader ------------------------------------------------------------------------------ + + +def test_parquet_reader_reads_schema_rows_and_exact_count(tmp_path): + _write_parquet(tmp_path / "d.parquet", PARQUET_ROWS) + result = get_reader("parquet").read(LocalFileSource(tmp_path), FileEntry("d.parquet", 0)) + + assert result.num_rows == 3 # exact, from the footer + assert result.rows_scanned == 3 + assert result.rows == PARQUET_ROWS + assert set(result.arrow_schema.names) == {"prompt", "score"} + + +def test_parquet_reader_row_cap_bounds_rows_but_keeps_exact_count(tmp_path): + _write_parquet(tmp_path / "d.parquet", PARQUET_ROWS) + result = get_reader("parquet").read(LocalFileSource(tmp_path), FileEntry("d.parquet", 0), row_cap=2) + + assert result.num_rows == 3 # footer count is unaffected by sampling + assert result.rows_scanned == 2 + assert result.rows == PARQUET_ROWS[:2] + + +def test_parquet_reader_zero_cap_reads_no_rows(tmp_path): + _write_parquet(tmp_path / "d.parquet", PARQUET_ROWS) + result = get_reader("parquet").read(LocalFileSource(tmp_path), FileEntry("d.parquet", 0), row_cap=0) + + assert result.rows == [] + assert result.num_rows == 3 + assert result.arrow_schema is not None # schema is still known without reading rows + + +# --- jsonl reader -------------------------------------------------------------------------------- + + +def test_jsonl_reader_full_read_is_exact_and_skips_blanks(tmp_path): + (tmp_path / "d.jsonl").write_text('{"a": 1}\n\n{"a": 2}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)) + + assert result.rows == [{"a": 1}, {"a": 2}] + assert result.num_rows == 2 # exact on a full read + assert result.arrow_schema is None # jsonl declares no schema + + +def test_jsonl_reader_row_cap_leaves_count_unknown(tmp_path): + (tmp_path / "d.jsonl").write_text('{"a": 1}\n{"a": 2}\n{"a": 3}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0), row_cap=2) + + assert result.rows == [{"a": 1}, {"a": 2}] + assert result.rows_scanned == 2 + assert result.num_rows is None # a partial read can't assert the total From c2dbd92a5e3162657aef5bc2e6d53f414198c905 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 15:21:27 -0400 Subject: [PATCH 03/44] feat(datasets): add profiling pipeline (partitions, splits, digest, envelope) Add the top-level profile() pipeline: list files, group into partitions by top-level directory, resolve splits by path inference (with canonical name normalization), read each file, and assemble a DatasetProfile with FileRecords, SplitProfiles, a listing content digest, and sampling metadata. Reads are exhaustive; per-file failures are isolated so one bad file never aborts a run. Row schema, column stats, and classification are stubbed pending later stages. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/__init__.py | 3 +- .../nemo_datasets_plugin/profiler/digest.py | 28 +++ .../profiler/partition.py | 39 ++++ .../nemo_datasets_plugin/profiler/pipeline.py | 115 ++++++++++++ .../nemo_datasets_plugin/profiler/splits.py | 72 ++++++++ plugins/nemo-datasets/tests/test_pipeline.py | 167 ++++++++++++++++++ 6 files changed, 423 insertions(+), 1 deletion(-) create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py create mode 100644 plugins/nemo-datasets/tests/test_pipeline.py diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py index 82bab560ba..f77643118b 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py @@ -9,5 +9,6 @@ """ from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource, LocalFileSource +from nemo_datasets_plugin.profiler.pipeline import profile -__all__ = ["FileEntry", "FileSource", "LocalFileSource"] +__all__ = ["FileEntry", "FileSource", "LocalFileSource", "profile"] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py new file mode 100644 index 0000000000..1f9a233795 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The content digest — a stable fingerprint of a dataset's file listing.""" + +from __future__ import annotations + +import hashlib + +from nemo_datasets_plugin.profiler.file_source import FileEntry + + +def content_digest(entries: list[FileEntry]) -> str: + """A stable digest over the (path, size, checksum) of every file, sorted by path. + + Uses only listing metadata — no file reads — so it is cheap and a profile can self-describe the + inputs it was built from; a mismatch on re-listing means the profile is stale. When a source + reports no checksum, (path, size) is the fallback, which cannot detect a same-size in-place edit. + """ + hasher = hashlib.sha256() + for entry in sorted(entries, key=lambda entry: entry.path): + hasher.update(entry.path.encode("utf-8")) + hasher.update(b"\0") + hasher.update(str(entry.size_bytes).encode("utf-8")) + hasher.update(b"\0") + hasher.update((entry.checksum or "").encode("utf-8")) + hasher.update(b"\n") + return f"sha256:{hasher.hexdigest()}" diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py new file mode 100644 index 0000000000..ab2b9557c9 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Partition grouping. + +A partition is a group of files profiled as a unit. This stage groups by top-level directory; a +later stage refines partitions whose files turn out to disagree on column schema. +""" + +from __future__ import annotations + +from pathlib import PurePosixPath + +from nemo_datasets_plugin.profiler.file_source import FileEntry + + +def _top_dir(path: str) -> str | None: + """The first path segment when the file is nested, else None for a root-level file.""" + parts = PurePosixPath(path).parts + return parts[0] if len(parts) > 1 else None + + +def group_partitions(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry]]]: + """Group files into (name, files) partitions by top-level directory. + + A single top-level group — every file at the root, or all under one container like ``data/`` — + is one "default" partition. Multiple top-level directories (e.g. ``main/`` and ``socratic/``) + each become their own partition, named after the directory; root-level files fall under + "default". + """ + by_dir: dict[str | None, list[FileEntry]] = {} + for entry in entries: + by_dir.setdefault(_top_dir(entry.path), []).append(entry) + + if len(by_dir) == 1: + return [("default", list(entries))] + + ordered = sorted(by_dir.items(), key=lambda item: (item[0] is None, item[0] or "")) + return [("default" if directory is None else directory, files) for directory, files in ordered] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py new file mode 100644 index 0000000000..4f4fb371db --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The top-level profiling pipeline. + +``profile(source)`` lists the files behind a :class:`FileSource`, groups them into partitions and +splits, reads them, and assembles a ``DatasetProfile``. This stage produces the structural envelope +— partitions, splits, FileRecords, content digest, and sampling metadata. The row schema, column +stats, and classification are added by later stages; until then each partition carries empty +``features`` / ``stats`` and an ``unknown`` classification. + +Reads are exhaustive (every row of every file). Sampling large datasets with bounded probes is a +later, drop-in optimization behind the same reader seam. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +from nemo_platform_plugin.files.dataset_profile import ( + DatasetProfile, + FileRecord, + PartitionClassification, + PartitionProfile, + SamplingInfo, + SplitProfile, +) + +from nemo_datasets_plugin.profiler.digest import content_digest +from nemo_datasets_plugin.profiler.file_source import FileSource +from nemo_datasets_plugin.profiler.partition import group_partitions +from nemo_datasets_plugin.profiler.readers import detect_format, get_reader +from nemo_datasets_plugin.profiler.splits import resolve_splits + +PROFILER_NAME = "nemo-dataset-profiler" +PROFILER_VERSION = "0.1.0" + + +def profile(source: FileSource, *, created_at: datetime | None = None) -> DatasetProfile: + """Profile the dataset behind ``source`` into a ``DatasetProfile``. + + ``created_at`` is injectable so a profile can be made reproducible byte-for-byte in tests; it + defaults to the current UTC time. + """ + created_at = created_at or datetime.now(timezone.utc) + all_entries = source.list_files() + data_entries = [entry for entry in all_entries if detect_format(entry.path) is not None] + + partitions: list[PartitionProfile] = [] + rows_scanned = 0 + all_exact = True + + for partition_name, partition_entries in group_partitions(data_entries): + split_profiles: list[SplitProfile] = [] + for split in resolve_splits(partition_entries): + file_records: list[FileRecord] = [] + split_examples = 0 + split_exact = True + for entry in split.entries: + reader = get_reader(detect_format(entry.path)) + try: + result = reader.read(source, entry) + num_rows = result.num_rows + rows_scanned += result.rows_scanned + except Exception: + # Failure isolation: keep the file's identity, skip its rows, keep going. + num_rows = None + file_records.append( + FileRecord( + path=entry.path, + size_bytes=entry.size_bytes, + checksum=entry.checksum, + num_rows=num_rows, + ) + ) + if num_rows is None: + split_exact = False + else: + split_examples += num_rows + all_exact = all_exact and split_exact + split_profiles.append( + SplitProfile( + name=split.name, + canonical=split.canonical, + files=file_records, + num_examples=split_examples if split_exact else None, + ) + ) + partitions.append( + PartitionProfile( + name=partition_name, + file_format=detect_format(partition_entries[0].path), + splits=split_profiles, + features=[], + stats={}, + classification=PartitionClassification(dataset_type="unknown"), + ) + ) + + sampling = SamplingInfo( + exhaustive=all_exact, + strategy="full", + rows_scanned=rows_scanned, + rows_total=rows_scanned if all_exact else None, + files_scanned=len(data_entries), + per_file_row_cap=None, + seed=None, + ) + return DatasetProfile( + content_digest=content_digest(all_entries), + created_at=created_at, + profiler_info={"name": PROFILER_NAME, "version": PROFILER_VERSION}, + sampling=sampling, + partitions=partitions, + ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py new file mode 100644 index 0000000000..ae582a0f20 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Split resolution from file paths. + +Given the files in one partition, group them into splits by inferring each file's split from its +name. A declared split map from a dataset card would take precedence over this inference, but card +parsing is not wired up yet, so path inference is the only source today. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import PurePosixPath + +from nemo_datasets_plugin.profiler.file_source import FileEntry + +# Strips a shard suffix like "-00000" or "-00000-of-00003" from a file stem. +_SHARD_SUFFIX = re.compile(r"-\d{2,}(?:-of-\d{2,})?$") + +# Common on-disk split words -> the canonical concept they normalize to. +_CANONICAL_ALIASES = { + "train": "train", + "test": "test", + "validation": "validation", + "valid": "validation", + "val": "validation", + "dev": "validation", +} + + +@dataclass(frozen=True) +class ResolvedSplit: + """A split and the files that belong to it.""" + + name: str # the on-disk split name, e.g. "train" or "train_prefs" + canonical: str | None # normalized concept (train | validation | test), or None + entries: list[FileEntry] + + +def _split_name(path: str) -> str: + """The shard-stripped file stem, e.g. train-00000-of-00003.parquet -> "train".""" + stem = PurePosixPath(path).name.split(".")[0] + return _SHARD_SUFFIX.sub("", stem) + + +def _canonical_for(split_name: str) -> str | None: + """Map a split name to its canonical concept, tolerating variant suffixes (train_prefs -> train).""" + lowered = split_name.lower() + for alias, canonical in _CANONICAL_ALIASES.items(): + if lowered == alias or lowered.startswith(f"{alias}_") or lowered.startswith(f"{alias}-"): + return canonical + return None + + +def resolve_splits(entries: list[FileEntry]) -> list[ResolvedSplit]: + """Group files into splits by path inference. + + Each file's split name is its shard-stripped stem; the canonical concept is matched against + common aliases (val/valid/dev -> validation). When no file carries a recognizable split, every + file lands in one "default" split. + """ + grouped: dict[str, list[FileEntry]] = {} + for entry in entries: + grouped.setdefault(_split_name(entry.path), []).append(entry) + + canonicals = {name: _canonical_for(name) for name in grouped} + if not any(canonicals.values()): + return [ResolvedSplit(name="default", canonical=None, entries=list(entries))] + + return [ResolvedSplit(name=name, canonical=canonicals[name], entries=grouped[name]) for name in sorted(grouped)] diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py new file mode 100644 index 0000000000..86956f515e --- /dev/null +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the profiling pipeline: digest, split/partition resolution, and envelope assembly.""" + +from datetime import datetime, timezone + +import pyarrow as pa +import pyarrow.parquet as pq +from nemo_datasets_plugin.profiler import profile +from nemo_datasets_plugin.profiler.digest import content_digest +from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource +from nemo_datasets_plugin.profiler.partition import group_partitions +from nemo_datasets_plugin.profiler.splits import resolve_splits + +FIXED_TIME = datetime(2026, 7, 13, 12, 0, 0, tzinfo=timezone.utc) + + +def _write_parquet(path, rows): + path.parent.mkdir(parents=True, exist_ok=True) + pq.write_table(pa.Table.from_pylist(rows), path) + + +def _entries(*paths): + return [FileEntry(path=p, size_bytes=100) for p in paths] + + +# --- content digest ------------------------------------------------------------------------------ + + +def test_content_digest_is_stable_and_order_independent(): + a = _entries("train.parquet", "test.parquet") + b = list(reversed(a)) + assert content_digest(a) == content_digest(b) + assert content_digest(a).startswith("sha256:") + + +def test_content_digest_changes_with_size(): + base = _entries("train.parquet") + bigger = [FileEntry(path="train.parquet", size_bytes=200)] + assert content_digest(base) != content_digest(bigger) + + +# --- split resolution ---------------------------------------------------------------------------- + + +def test_resolve_splits_infers_canonical_from_sharded_names(): + entries = _entries( + "train-00000-of-00002.parquet", + "train-00001-of-00002.parquet", + "validation-00000-of-00001.parquet", + ) + splits = {s.name: s for s in resolve_splits(entries)} + assert set(splits) == {"train", "validation"} + assert splits["train"].canonical == "train" + assert splits["validation"].canonical == "validation" + assert len(splits["train"].entries) == 2 + + +def test_resolve_splits_normalizes_aliases(): + splits = {s.name: s.canonical for s in resolve_splits(_entries("val.jsonl", "dev.jsonl"))} + assert splits == {"val": "validation", "dev": "validation"} + + +def test_resolve_splits_falls_back_to_single_default(): + splits = resolve_splits(_entries("shard-00000.parquet", "shard-00001.parquet")) + assert len(splits) == 1 + assert splits[0].name == "default" + assert splits[0].canonical is None + assert len(splits[0].entries) == 2 + + +# --- partition grouping -------------------------------------------------------------------------- + + +def test_group_partitions_single_default_for_root_files(): + assert group_partitions(_entries("train.parquet", "test.parquet")) == [ + ("default", _entries("train.parquet", "test.parquet")) + ] + + +def test_group_partitions_collapses_single_container_dir(): + parts = group_partitions(_entries("data/train.parquet", "data/test.parquet")) + assert [name for name, _ in parts] == ["default"] + + +def test_group_partitions_splits_multiple_top_dirs(): + parts = group_partitions(_entries("main/train.parquet", "socratic/train.parquet")) + assert [name for name, _ in parts] == ["main", "socratic"] + + +# --- end-to-end profile() ------------------------------------------------------------------------ + + +def test_profile_parquet_dataset_builds_envelope(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"prompt": "a"}, {"prompt": "b"}]) + _write_parquet(tmp_path / "validation-00000-of-00001.parquet", [{"prompt": "c"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.content_digest.startswith("sha256:") + assert result.profiler_info["name"] == "nemo-dataset-profiler" + assert len(result.partitions) == 1 + partition = result.partitions[0] + assert partition.name == "default" + assert partition.file_format == "parquet" + + splits = {s.name: s for s in partition.splits} + assert set(splits) == {"train", "validation"} + assert splits["train"].canonical == "train" + assert splits["train"].num_examples == 2 + assert splits["validation"].num_examples == 1 + assert splits["train"].files[0].num_rows == 2 + + # Structure envelope only: measurement and classification are stubbed for now. + assert partition.features == [] + assert partition.stats == {} + assert partition.classification.dataset_type == "unknown" + + assert result.sampling.exhaustive is True + assert result.sampling.strategy == "full" + assert result.sampling.rows_scanned == 3 + assert result.sampling.rows_total == 3 + assert result.sampling.files_scanned == 2 + + +def test_profile_jsonl_dataset_counts_rows_exactly(tmp_path): + (tmp_path / "train.jsonl").write_text('{"a": 1}\n{"a": 2}\n{"a": 3}\n') + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + partition = result.partitions[0] + assert partition.file_format == "jsonl" + assert partition.splits[0].name == "train" + assert partition.splits[0].num_examples == 3 + assert result.sampling.rows_scanned == 3 + + +def test_profile_multiple_directories_become_partitions(tmp_path): + _write_parquet(tmp_path / "main" / "train-00000-of-00001.parquet", [{"q": "1"}]) + _write_parquet(tmp_path / "socratic" / "train-00000-of-00001.parquet", [{"q": "2"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert [p.name for p in result.partitions] == ["main", "socratic"] + assert all(p.file_format == "parquet" for p in result.partitions) + + +def test_profile_isolates_unreadable_files(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + (tmp_path / "test-00000-of-00001.parquet").write_bytes(b"not a real parquet file") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + splits = {s.name: s for s in result.partitions[0].splits} + assert splits["train"].num_examples == 1 + assert splits["test"].num_examples is None # unreadable -> count unknown, not a crash + assert splits["test"].files[0].num_rows is None + assert result.sampling.exhaustive is False # a file could not be fully parsed + assert result.sampling.rows_total is None + + +def test_profile_is_deterministic(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) + source = LocalFileSource(tmp_path) + first = profile(source, created_at=FIXED_TIME) + second = profile(source, created_at=FIXED_TIME) + assert first.model_dump_json() == second.model_dump_json() From 84d4bc63fce8c4224d597c17b078af919fefe711 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 15:31:54 -0400 Subject: [PATCH 04/44] feat(datasets): derive row schema (features) from parquet and jsonl Add schema derivation: build the features tree de novo, converting parquet's declared arrow schema (keeping fixed-width numeric widths, recursing structs and lists, recording fixed_size_list lengths) and inferring from sampled rows for jsonl (dtype lattice with int/float widening, struct and list recursion). A list of {role, content} structs is recognized as the messages dtype. The pipeline now populates each partition's features; stats and classification remain stubbed. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/pipeline.py | 21 ++- .../nemo_datasets_plugin/profiler/schema.py | 140 ++++++++++++++++++ plugins/nemo-datasets/tests/test_pipeline.py | 5 +- plugins/nemo-datasets/tests/test_schema.py | 84 +++++++++++ 4 files changed, 242 insertions(+), 8 deletions(-) create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py create mode 100644 plugins/nemo-datasets/tests/test_schema.py diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 4f4fb371db..7406c7253d 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -5,9 +5,9 @@ ``profile(source)`` lists the files behind a :class:`FileSource`, groups them into partitions and splits, reads them, and assembles a ``DatasetProfile``. This stage produces the structural envelope -— partitions, splits, FileRecords, content digest, and sampling metadata. The row schema, column -stats, and classification are added by later stages; until then each partition carries empty -``features`` / ``stats`` and an ``unknown`` classification. +— partitions, splits, FileRecords, content digest, sampling metadata — and the derived row schema +(``features``). Column stats and classification are added by later stages; until then each partition +carries empty ``stats`` and an ``unknown`` classification. Reads are exhaustive (every row of every file). Sampling large datasets with bounded probes is a later, drop-in optimization behind the same reader seam. @@ -30,6 +30,7 @@ from nemo_datasets_plugin.profiler.file_source import FileSource from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.readers import detect_format, get_reader +from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import resolve_splits PROFILER_NAME = "nemo-dataset-profiler" @@ -51,6 +52,8 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase all_exact = True for partition_name, partition_entries in group_partitions(data_entries): + partition_rows: list[dict] = [] + arrow_schema = None split_profiles: list[SplitProfile] = [] for split in resolve_splits(partition_entries): file_records: list[FileRecord] = [] @@ -60,11 +63,17 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase reader = get_reader(detect_format(entry.path)) try: result = reader.read(source, entry) - num_rows = result.num_rows - rows_scanned += result.rows_scanned except Exception: # Failure isolation: keep the file's identity, skip its rows, keep going. + result = None + if result is None: num_rows = None + else: + num_rows = result.num_rows + rows_scanned += result.rows_scanned + partition_rows.extend(result.rows) + if arrow_schema is None: + arrow_schema = result.arrow_schema file_records.append( FileRecord( path=entry.path, @@ -91,7 +100,7 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase name=partition_name, file_format=detect_format(partition_entries[0].path), splits=split_profiles, - features=[], + features=derive_features(partition_rows, arrow_schema), stats={}, classification=PartitionClassification(dataset_type="unknown"), ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py new file mode 100644 index 0000000000..393798c4b3 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Row-schema derivation. + +Derive the ``features`` tree (a list of :class:`FeatureSchema`) de novo from the data. Parquet +carries a declared schema, so it is converted directly; formats without one (jsonl) are inferred +from the sampled rows by resolving each column's dtype. A list of ``{role, content}`` structs is +recognized as the ``messages`` dtype, and a list whose elements are all the same length records that +length as ``fixed_length``. +""" + +from __future__ import annotations + +from typing import Any + +import pyarrow as pa +from nemo_platform_plugin.files.dataset_profile import FeatureSchema + +# A list element carrying at least these keys is treated as a chat message (the messages dtype). +_MESSAGE_KEYS = {"role", "content"} + + +def derive_features(rows: list[dict[str, Any]], arrow_schema: pa.Schema | None = None) -> list[FeatureSchema]: + """The row schema. Uses the declared arrow schema when present, else infers from ``rows``.""" + if arrow_schema is not None: + return [ + _feature_from_arrow(arrow_schema.field(i).name, arrow_schema.field(i).type) + for i in range(len(arrow_schema)) + ] + return _features_from_rows(rows) + + +def _is_message_struct(item: FeatureSchema) -> bool: + return item.dtype == "struct" and item.fields is not None and _MESSAGE_KEYS <= {field.name for field in item.fields} + + +# --- from a declared arrow schema (parquet) ------------------------------------------------------ + +_ARROW_SCALAR_DTYPES = [ + (pa.types.is_boolean, "bool"), + (pa.types.is_int8, "int8"), + (pa.types.is_int16, "int16"), + (pa.types.is_int32, "int32"), + (pa.types.is_int64, "int64"), + (pa.types.is_uint8, "uint8"), + (pa.types.is_uint16, "uint16"), + (pa.types.is_uint32, "uint32"), + (pa.types.is_uint64, "uint64"), + (pa.types.is_float16, "float16"), + (pa.types.is_float32, "float32"), + (pa.types.is_float64, "float64"), + (pa.types.is_string, "string"), + (pa.types.is_large_string, "string"), +] + + +def _arrow_scalar_dtype(arrow_type: pa.DataType) -> str: + for predicate, dtype in _ARROW_SCALAR_DTYPES: + if predicate(arrow_type): + return dtype + return "json" + + +def _feature_from_arrow(name: str, arrow_type: pa.DataType) -> FeatureSchema: + if pa.types.is_struct(arrow_type): + fields = [ + _feature_from_arrow(arrow_type.field(i).name, arrow_type.field(i).type) + for i in range(arrow_type.num_fields) + ] + return FeatureSchema(name=name, dtype="struct", fields=fields) + if pa.types.is_fixed_size_list(arrow_type): + item = _feature_from_arrow("", arrow_type.value_type) + return FeatureSchema(name=name, dtype="list", items=item, fixed_length=arrow_type.list_size) + if pa.types.is_list(arrow_type) or pa.types.is_large_list(arrow_type): + item = _feature_from_arrow("", arrow_type.value_type) + dtype = "messages" if _is_message_struct(item) else "list" + return FeatureSchema(name=name, dtype=dtype, items=item) + return FeatureSchema(name=name, dtype=_arrow_scalar_dtype(arrow_type)) + + +# --- inferred from sampled rows (jsonl) ---------------------------------------------------------- + + +def _features_from_rows(rows: list[dict[str, Any]]) -> list[FeatureSchema]: + ordered_keys: list[str] = [] + seen: set[str] = set() + for row in rows: + for key in row: + if key not in seen: + seen.add(key) + ordered_keys.append(key) + return [_infer_feature(key, [row.get(key) for row in rows]) for key in ordered_keys] + + +def _infer_feature(name: str, values: list[Any]) -> FeatureSchema: + present = [value for value in values if value is not None] + if not present: + return FeatureSchema(name=name, dtype="json") + + if all(isinstance(value, dict) for value in present): + child_keys: list[str] = [] + seen: set[str] = set() + for record in present: + for key in record: + if key not in seen: + seen.add(key) + child_keys.append(key) + fields = [_infer_feature(key, [record.get(key) for record in present]) for key in child_keys] + return FeatureSchema(name=name, dtype="struct", fields=fields) + + if all(isinstance(value, list) for value in present): + item = _infer_feature("", [element for value in present for element in value]) + if _is_message_struct(item): + return FeatureSchema(name=name, dtype="messages", items=item) + lengths = {len(value) for value in present} + fixed_length = lengths.pop() if len(lengths) == 1 else None + return FeatureSchema(name=name, dtype="list", items=item, fixed_length=fixed_length) + + return FeatureSchema(name=name, dtype=_scalar_dtype(present)) + + +def _scalar_dtype(values: list[Any]) -> str: + dtypes: set[str] = set() + for value in values: + if isinstance(value, bool): # bool before int: bool is a subclass of int + dtypes.add("bool") + elif isinstance(value, int): + dtypes.add("int64") + elif isinstance(value, float): + dtypes.add("float64") + elif isinstance(value, str): + dtypes.add("string") + else: + dtypes.add("json") + if dtypes <= {"int64", "float64"} and dtypes: + return "float64" if "float64" in dtypes else "int64" + if len(dtypes) == 1: + return dtypes.pop() + return "json" diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 86956f515e..0c9ca495ae 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -112,8 +112,9 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert splits["validation"].num_examples == 1 assert splits["train"].files[0].num_rows == 2 - # Structure envelope only: measurement and classification are stubbed for now. - assert partition.features == [] + # Row schema is derived; stats and classification remain stubbed for now. + assert [f.name for f in partition.features] == ["prompt"] + assert partition.features[0].dtype == "string" assert partition.stats == {} assert partition.classification.dataset_type == "unknown" diff --git a/plugins/nemo-datasets/tests/test_schema.py b/plugins/nemo-datasets/tests/test_schema.py new file mode 100644 index 0000000000..71f3868b8a --- /dev/null +++ b/plugins/nemo-datasets/tests/test_schema.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for row-schema derivation (from a declared arrow schema and from sampled rows).""" + +import pyarrow as pa +from nemo_datasets_plugin.profiler.schema import derive_features + +# --- from a declared arrow schema (parquet) ------------------------------------------------------ + + +def test_from_arrow_scalars_keep_declared_widths(): + schema = pa.schema( + [("s", pa.string()), ("i", pa.int64()), ("i32", pa.int32()), ("b", pa.bool_()), ("f", pa.float64())] + ) + features = {f.name: f.dtype for f in derive_features([], schema)} + assert features == {"s": "string", "i": "int64", "i32": "int32", "b": "bool", "f": "float64"} + + +def test_from_arrow_list_of_role_content_structs_is_messages(): + schema = pa.schema([("prompt", pa.list_(pa.struct([("role", pa.string()), ("content", pa.string())])))]) + feature = derive_features([], schema)[0] + assert feature.dtype == "messages" + assert feature.items.dtype == "struct" + assert [f.name for f in feature.items.fields] == ["role", "content"] + + +def test_from_arrow_fixed_size_list_records_length(): + schema = pa.schema([("embedding", pa.list_(pa.float32(), 768))]) + feature = derive_features([], schema)[0] + assert feature.dtype == "list" + assert feature.fixed_length == 768 + assert feature.items.dtype == "float32" + + +def test_from_arrow_variable_list_has_no_fixed_length(): + feature = derive_features([], pa.schema([("tags", pa.list_(pa.string()))]))[0] + assert feature.dtype == "list" + assert feature.fixed_length is None + assert feature.items.dtype == "string" + + +# --- inferred from sampled rows (jsonl) ---------------------------------------------------------- + + +def test_from_rows_scalars_widen_int_and_float(): + rows = [{"a": 1, "b": 1.5, "c": "x", "d": True}, {"a": 2, "b": 2, "c": "y", "d": False}] + features = {f.name: f.dtype for f in derive_features(rows)} + assert features == {"a": "int64", "b": "float64", "c": "string", "d": "bool"} + + +def test_from_rows_list_of_role_content_structs_is_messages(): + rows = [{"conv": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]}] + feature = derive_features(rows)[0] + assert feature.dtype == "messages" + assert {f.name for f in feature.items.fields} == {"role", "content"} + + +def test_from_rows_constant_length_list_records_fixed_length(): + feature = derive_features([{"e": [0.1, 0.2, 0.3]}, {"e": [0.4, 0.5, 0.6]}])[0] + assert feature.dtype == "list" + assert feature.fixed_length == 3 + assert feature.items.dtype == "float64" + + +def test_from_rows_variable_length_list_has_no_fixed_length(): + feature = derive_features([{"e": [1, 2]}, {"e": [1, 2, 3]}])[0] + assert feature.dtype == "list" + assert feature.fixed_length is None + + +def test_from_rows_nested_struct(): + feature = derive_features([{"meta": {"id": 1, "src": "a"}}, {"meta": {"id": 2, "src": "b"}}])[0] + assert feature.dtype == "struct" + assert {f.name for f in feature.fields} == {"id", "src"} + + +def test_from_rows_all_null_column_is_json(): + assert derive_features([{"x": None}, {"x": None}])[0].dtype == "json" + + +def test_derive_features_prefers_declared_arrow_schema(): + feature = derive_features([{"x": 1}], pa.schema([("x", pa.int32())]))[0] + assert feature.dtype == "int32" # declared width beats the int64 inference from rows From 5b190b7c10e66f6e90067edd20f72633a051adee Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 15:58:47 -0400 Subject: [PATCH 05/44] feat(datasets): compute per-column stats Add per-column statistics keyed by dtype: length quantiles and cheap corruption signals for text, min/max/mean for numbers, chat-shape signals (turns, roles seen, ends-with-assistant, alternation, tool calls) for messages, and cardinality for both. Output is sparse; row values are never stored except a proven small enumeration under an exhaustive read. The pipeline now populates each partition's stats, leaving only classification. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/pipeline.py | 14 +- .../nemo_datasets_plugin/profiler/stats.py | 199 ++++++++++++++++++ plugins/nemo-datasets/tests/test_pipeline.py | 4 +- plugins/nemo-datasets/tests/test_stats.py | 101 +++++++++ 4 files changed, 311 insertions(+), 7 deletions(-) create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py create mode 100644 plugins/nemo-datasets/tests/test_stats.py diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 7406c7253d..ecdf6d5654 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -5,9 +5,9 @@ ``profile(source)`` lists the files behind a :class:`FileSource`, groups them into partitions and splits, reads them, and assembles a ``DatasetProfile``. This stage produces the structural envelope -— partitions, splits, FileRecords, content digest, sampling metadata — and the derived row schema -(``features``). Column stats and classification are added by later stages; until then each partition -carries empty ``stats`` and an ``unknown`` classification. +— partitions, splits, FileRecords, content digest, sampling metadata — the derived row schema +(``features``), and per-column ``stats``. Classification is added by a later stage; until then each +partition carries an ``unknown`` classification. Reads are exhaustive (every row of every file). Sampling large datasets with bounded probes is a later, drop-in optimization behind the same reader seam. @@ -32,6 +32,7 @@ from nemo_datasets_plugin.profiler.readers import detect_format, get_reader from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import resolve_splits +from nemo_datasets_plugin.profiler.stats import derive_stats PROFILER_NAME = "nemo-dataset-profiler" PROFILER_VERSION = "0.1.0" @@ -54,6 +55,7 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase for partition_name, partition_entries in group_partitions(data_entries): partition_rows: list[dict] = [] arrow_schema = None + partition_exact = True split_profiles: list[SplitProfile] = [] for split in resolve_splits(partition_entries): file_records: list[FileRecord] = [] @@ -87,6 +89,7 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase else: split_examples += num_rows all_exact = all_exact and split_exact + partition_exact = partition_exact and split_exact split_profiles.append( SplitProfile( name=split.name, @@ -95,13 +98,14 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase num_examples=split_examples if split_exact else None, ) ) + features = derive_features(partition_rows, arrow_schema) partitions.append( PartitionProfile( name=partition_name, file_format=detect_format(partition_entries[0].path), splits=split_profiles, - features=derive_features(partition_rows, arrow_schema), - stats={}, + features=features, + stats=derive_stats(features, partition_rows, exhaustive=partition_exact), classification=PartitionClassification(dataset_type="unknown"), ) ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py new file mode 100644 index 0000000000..b072dce32f --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-column statistics. + +Given a partition's features and its sampled rows, measure each top-level column according to its +dtype: length quantiles and corruption signals for text, min/max/mean for numbers, chat-shape +signals for messages, and cardinality for both. The result is sparse — a column with nothing worth +measuring is omitted. Row values themselves are never stored, except a proven small enumeration +under ``categorical.values`` when the read was exhaustive. +""" + +from __future__ import annotations + +import math +from typing import Any + +from nemo_platform_plugin.files.dataset_profile import ( + CategoricalStats, + ColumnStats, + FeatureSchema, + MessageStats, + NumericStats, + Quantiles, + TextQuality, + TextStats, +) + +# A proven enumeration is only stored when the read was exhaustive and this small. +_MAX_ENUM_VALUES = 32 + + +def derive_stats( + features: list[FeatureSchema], rows: list[dict[str, Any]], *, exhaustive: bool +) -> dict[str, ColumnStats]: + """Measure each top-level column. Keys are a subset of the feature names (sparse).""" + total = len(rows) + stats: dict[str, ColumnStats] = {} + for feature in features: + column = _column_stats(feature, [row.get(feature.name) for row in rows], total, exhaustive) + if column is not None: + stats[feature.name] = column + return stats + + +def _column_stats(feature: FeatureSchema, values: list[Any], total: int, exhaustive: bool) -> ColumnStats | None: + present = [value for value in values if value is not None] + null_rate = (total - len(present)) / total if total else 0.0 + + text = numeric = messages = categorical = quality = None + if feature.dtype == "string": + strings = [value for value in present if isinstance(value, str)] + if strings: + text = TextStats(chars=_quantiles([len(value) for value in strings])) + quality = _text_quality(strings) + counts = _cardinality(present, exhaustive) + if counts is not None and counts.distinct_count <= _MAX_ENUM_VALUES: + categorical = counts # a bounded string enumeration, not free text + elif _is_numeric(feature.dtype): + numbers = [float(value) for value in present if isinstance(value, (int, float)) and not isinstance(value, bool)] + if numbers: + numeric = NumericStats(min=min(numbers), max=max(numbers), mean=sum(numbers) / len(numbers)) + categorical = _cardinality(present, exhaustive) + elif feature.dtype == "messages": + messages = _message_stats([value for value in present if isinstance(value, list)]) + + column = ColumnStats( + null_rate=null_rate, text=text, numeric=numeric, messages=messages, categorical=categorical, quality=quality + ) + if not any([text, numeric, messages, categorical, quality]) and null_rate == 0.0: + return None # nothing worth measuring + return column + + +def _is_numeric(dtype: str) -> bool: + return dtype.startswith(("int", "uint", "float")) + + +def _quantiles(values: list[int]) -> Quantiles: + """Nearest-rank percentiles over the sample (n is small, so this stays exact).""" + ordered = sorted(values) + n = len(ordered) + + def at(percentile: int) -> int: + if n == 0: + return 0 + rank = math.ceil(percentile / 100 * n) + return ordered[min(rank, n) - 1] + + return Quantiles(p50=at(50), p95=at(95), p99=at(99), max=ordered[-1] if ordered else 0) + + +def _cardinality(present: list[Any], exhaustive: bool) -> CategoricalStats | None: + try: + distinct = set(present) + except TypeError: + return None # unhashable values (dicts / lists) have no cardinality signal + values = None + if exhaustive and len(distinct) <= _MAX_ENUM_VALUES: + values = sorted(str(value) for value in distinct) + return CategoricalStats(distinct_count=len(distinct), values=values) + + +# --- text quality -------------------------------------------------------------------------------- + + +def _text_quality(strings: list[str]) -> TextQuality: + total_chars = 0 + whitespace = 0 + non_ascii = 0 + repetition_sum = 0.0 + for value in strings: + total_chars += len(value) + whitespace += sum(char.isspace() for char in value) + non_ascii += sum(ord(char) > 127 for char in value) + repetition_sum += _repetition_score(value) + return TextQuality( + whitespace_ratio=whitespace / total_chars if total_chars else 0.0, + non_ascii_ratio=non_ascii / total_chars if total_chars else 0.0, + repetition_score=repetition_sum / len(strings) if strings else 0.0, + ) + + +def _repetition_score(text: str) -> float: + """Fraction of characters inside a run of the same character of length >= 4. + + A cheap corruption proxy: near zero for natural text, high for scraping junk and degenerate + single-character loops (``"aaaaaa"``, long ``"------"`` separators). + """ + if not text: + return 0.0 + redundant = 0 + run = 1 + for index in range(1, len(text)): + if text[index] == text[index - 1]: + run += 1 + else: + if run >= 4: + redundant += run + run = 1 + if run >= 4: + redundant += run + return redundant / len(text) + + +# --- messages ------------------------------------------------------------------------------------ + + +def _message_stats(rows_messages: list[list]) -> MessageStats | None: + if not rows_messages: + return None + turns: list[int] = [] + content_chars: list[int] = [] + roles_seen: list[str] = [] + ends_with_assistant = 0 + valid_alternation = 0 + has_tool_calls = False + + for messages in rows_messages: + turns.append(len(messages)) + total_content = 0 + for message in messages: + if not isinstance(message, dict): + continue + role = message.get("role") + if role is not None and role not in roles_seen: + roles_seen.append(role) + total_content += _content_len(message.get("content")) + if "tool_calls" in message or role == "tool": + has_tool_calls = True + content_chars.append(total_content) + if messages and isinstance(messages[-1], dict) and messages[-1].get("role") == "assistant": + ends_with_assistant += 1 + if _valid_alternation(messages): + valid_alternation += 1 + + n = len(rows_messages) + return MessageStats( + turns=_quantiles(turns), + content_chars=_quantiles(content_chars), + roles_seen=roles_seen, + ends_with_assistant_rate=ends_with_assistant / n, + valid_alternation_rate=valid_alternation / n, + has_tool_calls=has_tool_calls, + ) + + +def _content_len(content: Any) -> int: + if isinstance(content, str): + return len(content) + if isinstance(content, list): # VLM content as a list of typed parts + return sum(len(part.get("text", "")) for part in content if isinstance(part, dict)) + return 0 + + +def _valid_alternation(messages: list) -> bool: + """True when user/assistant turns alternate (ignoring any leading system turns).""" + roles = [m.get("role") for m in messages if isinstance(m, dict) and m.get("role") != "system"] + return all(roles[i] != roles[i + 1] for i in range(len(roles) - 1)) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 0c9ca495ae..c97b121653 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -112,10 +112,10 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert splits["validation"].num_examples == 1 assert splits["train"].files[0].num_rows == 2 - # Row schema is derived; stats and classification remain stubbed for now. + # Row schema and stats are derived; classification remains stubbed for now. assert [f.name for f in partition.features] == ["prompt"] assert partition.features[0].dtype == "string" - assert partition.stats == {} + assert partition.stats["prompt"].text is not None assert partition.classification.dataset_type == "unknown" assert result.sampling.exhaustive is True diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py new file mode 100644 index 0000000000..07a1a2cfba --- /dev/null +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for per-column statistics.""" + +from nemo_datasets_plugin.profiler.stats import derive_stats +from nemo_platform_plugin.files.dataset_profile import FeatureSchema + + +def _feature(name, dtype): + return FeatureSchema(name=name, dtype=dtype) + + +def _rows(name, values): + return [{name: value} for value in values] + + +# --- text ---------------------------------------------------------------------------------------- + + +def test_text_stats_length_quantiles_and_quality(): + values = ["a", "bb", "ccc", "dddd"] + stats = derive_stats([_feature("t", "string")], _rows("t", values), exhaustive=False)["t"] + assert stats.text.chars.max == 4 + assert stats.text.chars.p50 in {2, 3} # nearest-rank over 4 values + assert stats.quality is not None + assert stats.quality.whitespace_ratio == 0.0 + + +def test_text_quality_flags_repetition_and_non_ascii(): + stats = derive_stats([_feature("t", "string")], _rows("t", ["aaaaaaaa", "héllo wörld"]), exhaustive=False)["t"] + assert stats.quality.repetition_score > 0.0 # the "aaaaaaaa" run + assert stats.quality.non_ascii_ratio > 0.0 # accented characters + + +def test_free_text_string_has_no_categorical_but_low_cardinality_does(): + free_text = derive_stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)]), exhaustive=True) + assert free_text["t"].categorical is None # too many distinct values to be an enumeration + + labels = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"]), exhaustive=True) + assert labels["c"].categorical.distinct_count == 2 + assert labels["c"].categorical.values == ["no", "yes"] # proven enumeration under exhaustive read + + +# --- numeric ------------------------------------------------------------------------------------- + + +def test_numeric_stats_and_cardinality(): + stats = derive_stats([_feature("n", "int64")], _rows("n", [0, 4, 2, 2, 3]), exhaustive=True)["n"] + assert (stats.numeric.min, stats.numeric.max) == (0.0, 4.0) + assert stats.numeric.mean == 2.2 + assert stats.categorical.distinct_count == 4 # {0, 2, 3, 4} + + +def test_numeric_cardinality_values_withheld_when_not_exhaustive(): + stats = derive_stats([_feature("n", "int64")], _rows("n", [1, 2, 3]), exhaustive=False)["n"] + assert stats.categorical.distinct_count == 3 + assert stats.categorical.values is None # a sample cannot prove the enumeration + + +# --- messages ------------------------------------------------------------------------------------ + + +def test_message_stats_shape_signals(): + rows = [ + {"m": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello there"}]}, + {"m": [{"role": "user", "content": "again"}, {"role": "assistant", "content": "yes"}]}, + ] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.turns.max == 2 + assert stats.messages.roles_seen == ["user", "assistant"] # first-seen order + assert stats.messages.ends_with_assistant_rate == 1.0 + assert stats.messages.valid_alternation_rate == 1.0 + assert stats.messages.has_tool_calls is False + + +def test_message_stats_detects_tool_calls_and_user_ending(): + rows = [{"m": [{"role": "user", "content": "run"}, {"role": "assistant", "tool_calls": [{"id": "1"}]}]}] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.has_tool_calls is True + assert stats.messages.ends_with_assistant_rate == 1.0 # last turn is the assistant tool call + + +def test_message_ends_with_user_turn_is_prompt_only_signal(): + rows = [{"m": [{"role": "user", "content": "solve"}]}] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.ends_with_assistant_rate == 0.0 + + +# --- sparsity and null rate ---------------------------------------------------------------------- + + +def test_unmeasured_dtypes_are_omitted(): + features = [_feature("s", "struct"), _feature("j", "json")] + rows = [{"s": {"a": 1}, "j": object()}] + assert derive_stats(features, rows, exhaustive=False) == {} + + +def test_null_rate_is_reported(): + stats = derive_stats([_feature("t", "string")], _rows("t", ["a", None, "c", None]), exhaustive=False)["t"] + assert stats.null_rate == 0.5 From d1dd3dfe80d743182a85e37feba2f389745c4c4d Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 16:06:59 -0400 Subject: [PATCH 06/44] feat(datasets): classify roles, format/prompt-form axes, and dataset type Add classification: infer column roles from name aliases gated by dtype and stack them onto the feature nodes, resolve the format (standard/conversational/ mixed) and prompt-form (explicit/implicit/n-a) axes, and pick the most specific dataset type the assigned roles satisfy. The pipeline now emits a real classification per partition, replacing the unknown stub. Content probes and verifiability follow in a later stage. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/classify.py | 204 ++++++++++++++++++ .../nemo_datasets_plugin/profiler/pipeline.py | 11 +- plugins/nemo-datasets/tests/test_classify.py | 115 ++++++++++ plugins/nemo-datasets/tests/test_pipeline.py | 5 +- 4 files changed, 328 insertions(+), 7 deletions(-) create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py create mode 100644 plugins/nemo-datasets/tests/test_classify.py diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py new file mode 100644 index 0000000000..00606f7646 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Classification: assign column roles, resolve the format/prompt-form axes, and pick a dataset type. + +Roles are inferred from column names, gated by dtype, and stacked onto the feature nodes as +``semantic_role`` markers. The dataset type is the most specific structure the assigned roles +satisfy. Verifiability and content-probe corroboration are added by a later stage. +""" + +from __future__ import annotations + +from nemo_platform_plugin.files.dataset_profile import ( + ColumnStats, + Evidence, + FeatureSchema, + PartitionClassification, +) + +# Column-name aliases -> role. Score is handled separately (name alias + numeric dtype gate). +_ALIAS_ROLES = { + "prompt": "prompt", + "question": "prompt", + "instruction": "prompt", + "problem": "prompt", + "query": "prompt", + "context": "context", + "input": "context", + "passage": "context", + "document": "context", + "system": "system", + "system_prompt": "system", + "response": "completion", + "output": "completion", + "answer": "completion", + "completion": "completion", + "solution": "completion", + "messages": "messages", + "conversation": "messages", + "conversations": "messages", + "chosen": "chosen", + "rejected": "rejected", + "label": "label", + "rank": "rank", + "ground_truth": "ground_truth", + "reference_answer": "ground_truth", + "verification_info": "ground_truth", + "test_cases": "ground_truth", + "completions": "stepwise_completions", + "labels": "stepwise_labels", + "tools": "tools", + "image": "image", + "images": "image", + "id": "id", + "prompt_id": "id", + "source": "provenance", + "dataset": "provenance", + "model": "provenance", + "category": "meta", +} + +_SCORE_ALIASES = { + "score", + "score_chosen", + "score_rejected", + "helpfulness", + "correctness", + "coherence", + "complexity", + "verbosity", + "quality", + "rating", + "reward", +} + +_TEXT_DTYPES = {"string", "messages"} +# Roles whose string-vs-messages dtype decides the format axis. +_SHAPE_ROLES = {"prompt", "completion", "chosen", "rejected", "messages"} + + +def _is_numeric(dtype: str) -> bool: + return dtype.startswith(("int", "uint", "float")) + + +def _role_for(feature: FeatureSchema) -> str | None: + name = feature.name.lower() + dtype = feature.dtype + if name in _SCORE_ALIASES and _is_numeric(dtype): + return "score" + if name == "label" and dtype == "bool": + return "label" + + role = _ALIAS_ROLES.get(name) + if role is None: + return None + # dtype gates: reject an alias whose dtype contradicts the role. + if role == "messages" and dtype != "messages": + return None + if role in {"prompt", "completion", "chosen", "rejected", "context", "system", "ground_truth"}: + if dtype not in _TEXT_DTYPES: + return None + if role == "rank" and not _is_numeric(dtype): + return None + if role in {"stepwise_completions", "stepwise_labels"} and dtype != "list": + return None + if role == "label": # "label" reached here only when not bool + return None + return role + + +def _assign_roles(features: list[FeatureSchema]) -> None: + for feature in features: + role = _role_for(feature) + if role is not None: + feature.semantic_role = role + + +def _detect_modality(features: list[FeatureSchema]) -> str: + if any(feature.semantic_role == "image" or feature.dtype == "image" for feature in features): + return "image_text" + return "text" + + +def _detect_format(features: list[FeatureSchema]) -> str | None: + dtypes = {feature.dtype for feature in features if feature.semantic_role in _SHAPE_ROLES} + has_messages = "messages" in dtypes + has_string = "string" in dtypes + if has_messages and has_string: + return "mixed" + if has_messages: + return "conversational" + if has_string: + return "standard" + return None + + +def _detect_prompt_form(roles: set[str]) -> str | None: + if "prompt" in roles: + return "explicit" + if roles & {"chosen", "rejected", "completion"}: + return "implicit" # a prompt exists but is embedded in the completions + return "n/a" + + +def _messages_stats(features: list[FeatureSchema], stats: dict[str, ColumnStats]): + for feature in features: + if feature.semantic_role == "messages": + column = stats.get(feature.name) + if column is not None: + return column.messages + return None + + +def _detect_type(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> str: + roles = {feature.semantic_role for feature in features if feature.semantic_role} + + def has(*required: str) -> bool: + return all(role in roles for role in required) + + if has("prompt", "stepwise_completions", "stepwise_labels"): + return "stepwise_supervision" + if has("rank"): + return "ranked_responses" + if has("prompt", "completion", "score"): + return "scored_response" + if has("prompt", "completion", "label"): + return "unpaired_preference" + if has("chosen", "rejected"): + return "preference_pair" + if has("prompt", "completion"): + return "prompt_completion" + if "messages" in roles: + message_stats = _messages_stats(features, stats) + if message_stats is not None and message_stats.ends_with_assistant_rate < 0.5: + return "prompt_only" # a chat that ends on a user turn has no training target + return "messages" + if "prompt" in roles: + return "prompt_only" + if len(features) == 1 and features[0].dtype == "string" and features[0].semantic_role is None: + return "text" + return "unknown" + + +def classify(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> PartitionClassification: + """Assign roles onto ``features`` in place and return the partition's classification.""" + _assign_roles(features) + roles = {feature.semantic_role for feature in features if feature.semantic_role} + dataset_type = _detect_type(features, stats) + fmt = _detect_format(features) + + evidence: list[Evidence] = [] + role_columns = [f"{feature.name} -> {feature.semantic_role}" for feature in features if feature.semantic_role] + if role_columns: + evidence.append(Evidence(kind="column_name", detail=f"columns matched roles: {', '.join(role_columns)}")) + if fmt is not None: + evidence.append(Evidence(kind="column_dtype", detail=f"{fmt} format from role column dtypes")) + + return PartitionClassification( + modality=_detect_modality(features), + dataset_type=dataset_type, + format=fmt, + prompt_form=_detect_prompt_form(roles) if dataset_type != "unknown" else None, + evidence=evidence, + ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index ecdf6d5654..27f54e8597 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -6,8 +6,8 @@ ``profile(source)`` lists the files behind a :class:`FileSource`, groups them into partitions and splits, reads them, and assembles a ``DatasetProfile``. This stage produces the structural envelope — partitions, splits, FileRecords, content digest, sampling metadata — the derived row schema -(``features``), and per-column ``stats``. Classification is added by a later stage; until then each -partition carries an ``unknown`` classification. +(``features``), per-column ``stats``, and a ``classification`` (roles, format, prompt form, and +dataset type). Verifiability and content-probe corroboration are added by a later stage. Reads are exhaustive (every row of every file). Sampling large datasets with bounded probes is a later, drop-in optimization behind the same reader seam. @@ -20,12 +20,12 @@ from nemo_platform_plugin.files.dataset_profile import ( DatasetProfile, FileRecord, - PartitionClassification, PartitionProfile, SamplingInfo, SplitProfile, ) +from nemo_datasets_plugin.profiler.classify import classify from nemo_datasets_plugin.profiler.digest import content_digest from nemo_datasets_plugin.profiler.file_source import FileSource from nemo_datasets_plugin.profiler.partition import group_partitions @@ -99,14 +99,15 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase ) ) features = derive_features(partition_rows, arrow_schema) + stats = derive_stats(features, partition_rows, exhaustive=partition_exact) partitions.append( PartitionProfile( name=partition_name, file_format=detect_format(partition_entries[0].path), splits=split_profiles, features=features, - stats=derive_stats(features, partition_rows, exhaustive=partition_exact), - classification=PartitionClassification(dataset_type="unknown"), + stats=stats, + classification=classify(features, stats), ) ) diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py new file mode 100644 index 0000000000..94795c0355 --- /dev/null +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for classification: role assignment, format/prompt-form axes, and dataset type.""" + +from nemo_datasets_plugin.profiler.classify import classify +from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema, MessageStats, Quantiles + + +def _f(name, dtype): + return FeatureSchema(name=name, dtype=dtype) + + +def _messages_column(ends_with_assistant_rate): + q = Quantiles(p50=1, p95=1, p99=1, max=1) + return ColumnStats( + messages=MessageStats( + turns=q, + content_chars=q, + roles_seen=["user", "assistant"], + ends_with_assistant_rate=ends_with_assistant_rate, + valid_alternation_rate=1.0, + ) + ) + + +# --- roles --------------------------------------------------------------------------------------- + + +def test_roles_assigned_by_name_and_dtype(): + features = [_f("prompt", "string"), _f("response", "string"), _f("helpfulness", "int64")] + classify(features, {}) + assert [f.semantic_role for f in features] == ["prompt", "completion", "score"] + + +def test_dtype_gate_rejects_mismatched_aliases(): + # "label" only counts as a label when boolean; a string column named "messages" is not messages. + features = [_f("label", "string"), _f("messages", "string")] + classify(features, {}) + assert all(f.semantic_role is None for f in features) + + +def test_physical_name_differs_from_role(): + features = [_f("response", "string")] + classify(features, {}) + assert features[0].semantic_role == "completion" + + +# --- format axis --------------------------------------------------------------------------------- + + +def test_format_standard_conversational_and_mixed(): + assert classify([_f("prompt", "string"), _f("completion", "string")], {}).format == "standard" + assert classify([_f("prompt", "messages"), _f("completion", "messages")], {}).format == "conversational" + mixed = [_f("prompt", "string"), _f("chosen", "messages"), _f("rejected", "messages")] + assert classify(mixed, {}).format == "mixed" + + +# --- dataset type + prompt form ------------------------------------------------------------------ + + +def test_prompt_completion_with_explicit_prompt(): + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert result.dataset_type == "prompt_completion" + assert result.prompt_form == "explicit" + + +def test_preference_pair_is_implicit_without_a_prompt(): + result = classify([_f("chosen", "string"), _f("rejected", "string")], {}) + assert result.dataset_type == "preference_pair" + assert result.prompt_form == "implicit" + + +def test_scored_response_beats_prompt_completion(): + features = [ + _f("prompt", "string"), + _f("response", "string"), + _f("helpfulness", "int64"), + _f("correctness", "int64"), + ] + assert classify(features, {}).dataset_type == "scored_response" + + +def test_unpaired_preference_needs_boolean_label(): + features = [_f("prompt", "string"), _f("completion", "string"), _f("label", "bool")] + assert classify(features, {}).dataset_type == "unpaired_preference" + + +def test_messages_ending_on_assistant_is_messages_type(): + result = classify([_f("messages", "messages")], {"messages": _messages_column(1.0)}) + assert result.dataset_type == "messages" + assert result.prompt_form == "n/a" + + +def test_messages_ending_on_user_is_prompt_only(): + result = classify([_f("messages", "messages")], {"messages": _messages_column(0.0)}) + assert result.dataset_type == "prompt_only" + + +def test_single_text_column_is_text(): + assert classify([_f("text", "string")], {}).dataset_type == "text" + + +def test_unrecognized_columns_are_unknown(): + result = classify([_f("foo", "int64"), _f("bar", "int64")], {}) + assert result.dataset_type == "unknown" + assert result.prompt_form is None # no axes asserted for unknown data + + +# --- evidence ------------------------------------------------------------------------------------ + + +def test_classification_records_evidence(): + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert {e.kind for e in result.evidence} >= {"column_name", "column_dtype"} diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index c97b121653..12d4f7fa54 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -112,11 +112,12 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert splits["validation"].num_examples == 1 assert splits["train"].files[0].num_rows == 2 - # Row schema and stats are derived; classification remains stubbed for now. + # Row schema, stats, and classification are all derived now. assert [f.name for f in partition.features] == ["prompt"] assert partition.features[0].dtype == "string" + assert partition.features[0].semantic_role == "prompt" assert partition.stats["prompt"].text is not None - assert partition.classification.dataset_type == "unknown" + assert partition.classification.dataset_type == "prompt_only" # a lone prompt column, no target assert result.sampling.exhaustive is True assert result.sampling.strategy == "full" From 1f194d0e93c15084d2574f81cf177e102b4d091a Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 16:21:31 -0400 Subject: [PATCH 07/44] feat(datasets): detect verifiability and implicit-prompt probes Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/classify.py | 113 +++++++++++++++++- .../nemo_datasets_plugin/profiler/pipeline.py | 6 +- plugins/nemo-datasets/tests/test_classify.py | 40 +++++++ plugins/nemo-datasets/tests/test_pipeline.py | 13 ++ 4 files changed, 166 insertions(+), 6 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index 00606f7646..1daafa7545 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -10,11 +10,14 @@ from __future__ import annotations +import re + from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, Evidence, FeatureSchema, PartitionClassification, + Verifiability, ) # Column-name aliases -> role. Score is handled separately (name alias + numeric dtype gate). @@ -181,12 +184,111 @@ def has(*required: str) -> bool: return "unknown" -def classify(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> PartitionClassification: - """Assign roles onto ``features`` in place and return the partition's classification.""" +# --- content probes ------------------------------------------------------------------------------ + +_TRANSCRIPT_MARKER = re.compile(r"\n\n(?:Human|Assistant|User):") +_GSM8K_ANSWER = re.compile(r"####\s*-?[\d.,/]+\s*$") +_BOXED_ANSWER = re.compile(r"\\boxed\{") + + +def _pct(fraction: float) -> str: + return f"{round(fraction * 100)}%" + + +def _completion_texts(features: list[FeatureSchema], rows: list[dict]) -> list[str]: + completion = next((feature for feature in features if feature.semantic_role == "completion"), None) + if completion is None: + return [] + texts: list[str] = [] + for row in rows: + value = row.get(completion.name) + if isinstance(value, str): + texts.append(value) + elif isinstance(value, list) and value and isinstance(value[-1], dict): + content = value[-1].get("content") # the final assistant turn for a conversational completion + if isinstance(content, str): + texts.append(content) + return texts + + +def _detect_verifiability(features: list[FeatureSchema], rows: list[dict]) -> Verifiability | None: + if not rows: + return None + + ground_truth = next((feature for feature in features if feature.semantic_role == "ground_truth"), None) + if ground_truth is not None: + present = sum(1 for row in rows if row.get(ground_truth.name) not in (None, "", [])) + coverage = present / len(rows) + detail = f"'{ground_truth.name}' present in {_pct(coverage)} of {len(rows)} sampled rows" + return Verifiability( + method="ground_truth_column", coverage=coverage, evidence=[Evidence(kind="content_probe", detail=detail)] + ) + + texts = _completion_texts(features, rows) + if texts: + hits = sum(1 for text in texts if _GSM8K_ANSWER.search(text) or _BOXED_ANSWER.search(text)) + if hits: + coverage = hits / len(texts) + detail = f"completion ends with an extractable answer (#### or \\boxed) in {_pct(coverage)} of {len(texts)} sampled rows" + return Verifiability( + method="extractable_final_answer", + coverage=coverage, + evidence=[Evidence(kind="content_probe", detail=detail)], + ) + return None + + +def _common_prefix_len(left: str, right: str) -> int: + limit = min(len(left), len(right)) + index = 0 + while index < limit and left[index] == right[index]: + index += 1 + return index + + +def _implicit_prompt_evidence(features: list[FeatureSchema], rows: list[dict]) -> Evidence | None: + if not rows: + return None + + targets = [f for f in features if f.semantic_role in {"chosen", "rejected", "completion"} and f.dtype == "string"] + texts = [row.get(f.name) for f in targets for row in rows if isinstance(row.get(f.name), str)] + if texts: + marked = sum(1 for text in texts if _TRANSCRIPT_MARKER.search(text)) + if marked: + detail = f"embedded transcript markers in {_pct(marked / len(texts))} of sampled completions - prompt is embedded" + return Evidence(kind="content_probe", detail=detail) + + chosen = next((f for f in features if f.semantic_role == "chosen" and f.dtype == "string"), None) + rejected = next((f for f in features if f.semantic_role == "rejected" and f.dtype == "string"), None) + if chosen is not None and rejected is not None: + pairs = 0 + shared = 0 + for row in rows: + left, right = row.get(chosen.name), row.get(rejected.name) + if isinstance(left, str) and isinstance(right, str): + pairs += 1 + if _common_prefix_len(left, right) >= 16: + shared += 1 + if pairs and shared / pairs >= 0.5: + detail = f"chosen/rejected share a common prefix in {_pct(shared / pairs)} of pairs - prompt is embedded" + return Evidence(kind="content_probe", detail=detail) + return None + + +def classify( + features: list[FeatureSchema], stats: dict[str, ColumnStats], rows: list[dict] | None = None +) -> PartitionClassification: + """Assign roles onto ``features`` in place and return the partition's classification. + + ``rows`` (the sampled rows) drive the content probes — verifiability and implicit-prompt + detection; role/axis/type inference needs only the schema and stats. + """ + rows = rows or [] _assign_roles(features) roles = {feature.semantic_role for feature in features if feature.semantic_role} dataset_type = _detect_type(features, stats) fmt = _detect_format(features) + prompt_form = _detect_prompt_form(roles) if dataset_type != "unknown" else None evidence: list[Evidence] = [] role_columns = [f"{feature.name} -> {feature.semantic_role}" for feature in features if feature.semantic_role] @@ -194,11 +296,16 @@ def classify(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> Pa evidence.append(Evidence(kind="column_name", detail=f"columns matched roles: {', '.join(role_columns)}")) if fmt is not None: evidence.append(Evidence(kind="column_dtype", detail=f"{fmt} format from role column dtypes")) + if prompt_form == "implicit": + embedded = _implicit_prompt_evidence(features, rows) + if embedded is not None: + evidence.append(embedded) return PartitionClassification( modality=_detect_modality(features), dataset_type=dataset_type, format=fmt, - prompt_form=_detect_prompt_form(roles) if dataset_type != "unknown" else None, + prompt_form=prompt_form, + verifiability=_detect_verifiability(features, rows), evidence=evidence, ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 27f54e8597..ff04dd065e 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -6,8 +6,8 @@ ``profile(source)`` lists the files behind a :class:`FileSource`, groups them into partitions and splits, reads them, and assembles a ``DatasetProfile``. This stage produces the structural envelope — partitions, splits, FileRecords, content digest, sampling metadata — the derived row schema -(``features``), per-column ``stats``, and a ``classification`` (roles, format, prompt form, and -dataset type). Verifiability and content-probe corroboration are added by a later stage. +(``features``), per-column ``stats``, and the full ``classification`` (roles, format, prompt form, +dataset type, and verifiability). Reads are exhaustive (every row of every file). Sampling large datasets with bounded probes is a later, drop-in optimization behind the same reader seam. @@ -107,7 +107,7 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase splits=split_profiles, features=features, stats=stats, - classification=classify(features, stats), + classification=classify(features, stats, partition_rows), ) ) diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index 94795c0355..235be47945 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -113,3 +113,43 @@ def test_unrecognized_columns_are_unknown(): def test_classification_records_evidence(): result = classify([_f("prompt", "string"), _f("completion", "string")], {}) assert {e.kind for e in result.evidence} >= {"column_name", "column_dtype"} + + +# --- verifiability + content probes -------------------------------------------------------------- + + +def test_verifiability_extractable_gsm8k_answer(): + features = [_f("problem", "string"), _f("solution", "string")] + rows = [{"problem": "q", "solution": "steps #### 18"}, {"problem": "q", "solution": "no final answer"}] + result = classify(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 0.5 + + +def test_verifiability_boxed_answer(): + features = [_f("prompt", "string"), _f("completion", "string")] + result = classify(features, {}, [{"prompt": "q", "completion": r"reasoning \boxed{42}"}]) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + + +def test_verifiability_ground_truth_column_coverage(): + features = [_f("prompt", "string"), _f("ground_truth", "string")] + rows = [{"prompt": "q", "ground_truth": "42"}, {"prompt": "q", "ground_truth": None}] + result = classify(features, {}, rows) + assert result.verifiability.method == "ground_truth_column" + assert result.verifiability.coverage == 0.5 + + +def test_no_verifiability_without_a_target(): + features = [_f("prompt", "string"), _f("completion", "string")] + result = classify(features, {}, [{"prompt": "q", "completion": "just prose, no answer"}]) + assert result.verifiability is None + + +def test_implicit_prompt_evidence_from_embedded_transcript(): + features = [_f("chosen", "string"), _f("rejected", "string")] + rows = [{"chosen": "\n\nHuman: hi\n\nAssistant: hello", "rejected": "\n\nHuman: hi\n\nAssistant: hey"}] + result = classify(features, {}, rows) + assert result.prompt_form == "implicit" + assert any(e.kind == "content_probe" for e in result.evidence) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 12d4f7fa54..2cded58b87 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -161,6 +161,19 @@ def test_profile_isolates_unreadable_files(tmp_path): assert result.sampling.rows_total is None +def test_profile_classifies_roles_type_and_verifiability(tmp_path): + _write_parquet( + tmp_path / "train-00000-of-00001.parquet", + [{"problem": "q1", "solution": "steps #### 5"}, {"problem": "q2", "solution": "steps #### 6"}], + ) + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert {f.semantic_role for f in partition.features} == {"prompt", "completion"} + assert partition.classification.dataset_type == "prompt_completion" + assert partition.classification.verifiability.method == "extractable_final_answer" + assert partition.classification.verifiability.coverage == 1.0 + + def test_profile_is_deterministic(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) source = LocalFileSource(tmp_path) From 261174a44d2a751490701913e7954eeacb24ece2 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Mon, 13 Jul 2026 17:15:06 -0400 Subject: [PATCH 08/44] fix: harden the profiler and wire up the profile CLI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the correctness issues found reviewing the dataset-profile-engine, and bring the plugin in line with the repo's packaging conventions. Correctness: - readers/jsonl: skip lines that parse to non-objects, so a stray scalar or array no longer poisons the (unprotected) schema/stats stage and aborts the whole profile — the reader now honors its list[dict] contract. - stats: _content_len tolerates a present-but-non-string `text` in a VLM content part instead of raising len(None). - pipeline: digest over the stored FileRecords (data files) rather than every listed file, so the digest is recomputable from the profile and a non-data file (e.g. a README) no longer flips it. - cli: `nemo datasets profile` now runs the engine and prints JSON or YAML (validating --output), replacing the not-implemented stub. - classify: allow `ground_truth` (and its test_cases/verification_info aliases) to carry list/struct/json dtypes so container verification targets are detected; a bare numeric column is still rejected. - pipeline: distinguish "exact row count known" (footer) from "every row scanned" so `exhaustive` stays correct once bounded sampling lands. - pipeline: move the reader lookup inside failure isolation, so a detected format with no registered reader is skipped per-file, not fatal. Packaging & hygiene: - Drop the three __init__.py files for a PEP 420 namespace package, matching the other plugins; built-in readers self-register lazily on the first get_reader() call (no import cycle, no pyarrow at import time). - Resolve ty diagnostics via a _format_of() helper and walrus narrowing. Add 11 regression tests covering each fix. Signed-off-by: Albert Cui --- .../src/nemo_datasets_plugin/__init__.py | 8 --- .../src/nemo_datasets_plugin/cli.py | 19 +++++- .../nemo_datasets_plugin/profiler/__init__.py | 14 ----- .../nemo_datasets_plugin/profiler/classify.py | 9 ++- .../nemo_datasets_plugin/profiler/pipeline.py | 63 ++++++++++++------- .../profiler/readers/__init__.py | 28 --------- .../profiler/readers/base.py | 20 +++++- .../profiler/readers/jsonl.py | 5 +- .../profiler/readers/parquet.py | 1 - .../nemo_datasets_plugin/profiler/stats.py | 4 +- plugins/nemo-datasets/tests/test_classify.py | 24 +++++++ plugins/nemo-datasets/tests/test_cli.py | 37 +++++++++++ plugins/nemo-datasets/tests/test_pipeline.py | 47 +++++++++++++- plugins/nemo-datasets/tests/test_readers.py | 11 +++- plugins/nemo-datasets/tests/test_stats.py | 7 +++ 15 files changed, 212 insertions(+), 85 deletions(-) delete mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py delete mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py delete mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py deleted file mode 100644 index ecbf54d3aa..0000000000 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Dataset profiler for NeMo Platform filesets. - -Computes a ``DatasetProfile`` (the contract in ``nemo_platform_plugin.files.dataset_profile``) from -a directory of dataset files, and exposes it as the ``nemo datasets`` CLI. -""" diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py index 0dcbc27343..497b7eecbb 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py @@ -24,7 +24,22 @@ def profile( output: str = typer.Option("json", "--output", "-o", help="Output format: json | yaml."), ) -> None: """Profile a local dataset directory and print its DatasetProfile.""" - # The profiling core lands in a follow-up commit; this wires up the command surface. - raise typer.BadParameter("dataset profiling is not implemented yet") + from nemo_datasets_plugin.profiler.file_source import LocalFileSource + from nemo_datasets_plugin.profiler.pipeline import profile as run_profile + + if output not in {"json", "yaml"}: + raise typer.BadParameter("output must be 'json' or 'yaml'") + try: + source = LocalFileSource(path) + except NotADirectoryError as exc: + raise typer.BadParameter(str(exc)) from exc + + result = run_profile(source) + if output == "yaml": + import yaml + + typer.echo(yaml.safe_dump(result.model_dump(mode="json"), sort_keys=False)) + else: + typer.echo(result.model_dump_json(indent=2)) return app diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py deleted file mode 100644 index f77643118b..0000000000 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""The dataset profiler core. - -A dependency-free library (pyarrow + stdlib) that reads dataset files through a -:class:`~nemo_datasets_plugin.profiler.file_source.FileSource` seam and per-format readers, and — in -later commits — measures and classifies them into a ``DatasetProfile``. -""" - -from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource, LocalFileSource -from nemo_datasets_plugin.profiler.pipeline import profile - -__all__ = ["FileEntry", "FileSource", "LocalFileSource", "profile"] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index 1daafa7545..ef29e2f036 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -77,6 +77,9 @@ } _TEXT_DTYPES = {"string", "messages"} +# A verification target is naturally a container: test_cases (list), verification_info (struct), or a +# plain string answer — but never a bare scalar/number, which is far more likely a label or score. +_GROUND_TRUTH_DTYPES = {"string", "messages", "list", "struct", "json"} # Roles whose string-vs-messages dtype decides the format axis. _SHAPE_ROLES = {"prompt", "completion", "chosen", "rejected", "messages"} @@ -99,9 +102,11 @@ def _role_for(feature: FeatureSchema) -> str | None: # dtype gates: reject an alias whose dtype contradicts the role. if role == "messages" and dtype != "messages": return None - if role in {"prompt", "completion", "chosen", "rejected", "context", "system", "ground_truth"}: + if role in {"prompt", "completion", "chosen", "rejected", "context", "system"}: if dtype not in _TEXT_DTYPES: return None + if role == "ground_truth" and dtype not in _GROUND_TRUTH_DTYPES: + return None if role == "rank" and not _is_numeric(dtype): return None if role in {"stepwise_completions", "stepwise_labels"} and dtype != "list": @@ -251,7 +256,7 @@ def _implicit_prompt_evidence(features: list[FeatureSchema], rows: list[dict]) - return None targets = [f for f in features if f.semantic_role in {"chosen", "rejected", "completion"} and f.dtype == "string"] - texts = [row.get(f.name) for f in targets for row in rows if isinstance(row.get(f.name), str)] + texts = [value for f in targets for row in rows if isinstance((value := row.get(f.name)), str)] if texts: marked = sum(1 for text in texts if _TRANSCRIPT_MARKER.search(text)) if marked: diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index ff04dd065e..683ee9e70d 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -17,27 +17,35 @@ from datetime import datetime, timezone -from nemo_platform_plugin.files.dataset_profile import ( - DatasetProfile, - FileRecord, - PartitionProfile, - SamplingInfo, - SplitProfile, -) - from nemo_datasets_plugin.profiler.classify import classify from nemo_datasets_plugin.profiler.digest import content_digest from nemo_datasets_plugin.profiler.file_source import FileSource from nemo_datasets_plugin.profiler.partition import group_partitions -from nemo_datasets_plugin.profiler.readers import detect_format, get_reader +from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import resolve_splits from nemo_datasets_plugin.profiler.stats import derive_stats +from nemo_platform_plugin.files.dataset_profile import ( + DatasetProfile, + FileRecord, + PartitionProfile, + SamplingInfo, + SplitProfile, +) PROFILER_NAME = "nemo-dataset-profiler" PROFILER_VERSION = "0.1.0" +def _format_of(path: str) -> str: + """The registered format of a data file. Callers pass only pre-filtered ``data_entries``, so the + format is always known; a None here would mean that invariant was broken.""" + file_format = detect_format(path) + if file_format is None: + raise ValueError(f"no registered format for {path!r}") + return file_format + + def profile(source: FileSource, *, created_at: datetime | None = None) -> DatasetProfile: """Profile the dataset behind ``source`` into a ``DatasetProfile``. @@ -50,32 +58,36 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase partitions: list[PartitionProfile] = [] rows_scanned = 0 - all_exact = True + all_scanned = True for partition_name, partition_entries in group_partitions(data_entries): partition_rows: list[dict] = [] arrow_schema = None - partition_exact = True + partition_scanned = True split_profiles: list[SplitProfile] = [] for split in resolve_splits(partition_entries): file_records: list[FileRecord] = [] split_examples = 0 - split_exact = True + split_counts_known = True # every file's exact total row count is known (footer or full scan) + split_scanned = True # every row of every file was actually parsed for entry in split.entries: - reader = get_reader(detect_format(entry.path)) try: - result = reader.read(source, entry) + result = get_reader(_format_of(entry.path)).read(source, entry) except Exception: - # Failure isolation: keep the file's identity, skip its rows, keep going. + # Failure isolation: an unreadable file (or missing reader) keeps its identity, + # skips its rows, and does not abort the profile. result = None if result is None: num_rows = None + scanned_all = False else: num_rows = result.num_rows rows_scanned += result.rows_scanned partition_rows.extend(result.rows) if arrow_schema is None: arrow_schema = result.arrow_schema + # Exhaustive requires parsing every row; a known footer count alone is not enough. + scanned_all = num_rows is not None and result.rows_scanned >= num_rows file_records.append( FileRecord( path=entry.path, @@ -85,25 +97,27 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase ) ) if num_rows is None: - split_exact = False + split_counts_known = False else: split_examples += num_rows - all_exact = all_exact and split_exact - partition_exact = partition_exact and split_exact + if not scanned_all: + split_scanned = False + all_scanned = all_scanned and split_scanned + partition_scanned = partition_scanned and split_scanned split_profiles.append( SplitProfile( name=split.name, canonical=split.canonical, files=file_records, - num_examples=split_examples if split_exact else None, + num_examples=split_examples if split_counts_known else None, ) ) features = derive_features(partition_rows, arrow_schema) - stats = derive_stats(features, partition_rows, exhaustive=partition_exact) + stats = derive_stats(features, partition_rows, exhaustive=partition_scanned) partitions.append( PartitionProfile( name=partition_name, - file_format=detect_format(partition_entries[0].path), + file_format=_format_of(partition_entries[0].path), splits=split_profiles, features=features, stats=stats, @@ -112,16 +126,17 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase ) sampling = SamplingInfo( - exhaustive=all_exact, + exhaustive=all_scanned, strategy="full", rows_scanned=rows_scanned, - rows_total=rows_scanned if all_exact else None, + rows_total=rows_scanned if all_scanned else None, files_scanned=len(data_entries), per_file_row_cap=None, seed=None, ) return DatasetProfile( - content_digest=content_digest(all_entries), + # Digest only the files stored as FileRecords, so the profile can recompute its own digest. + content_digest=content_digest(data_entries), created_at=created_at, profiler_info={"name": PROFILER_NAME, "version": PROFILER_VERSION}, sampling=sampling, diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py deleted file mode 100644 index 94ad5b6e52..0000000000 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Format readers, registered by importing this package. - -Importing the package populates the registry (each reader module self-registers), so -:func:`get_reader` resolves every supported format. -""" - -from nemo_datasets_plugin.profiler.readers.base import ( - FormatReader, - ReadResult, - detect_format, - get_reader, - register_reader, -) - -# Import for side effects: each module calls register_reader() at import time. -from nemo_datasets_plugin.profiler.readers import jsonl as _jsonl # noqa: F401 isort:skip -from nemo_datasets_plugin.profiler.readers import parquet as _parquet # noqa: F401 isort:skip - -__all__ = [ - "FormatReader", - "ReadResult", - "detect_format", - "get_reader", - "register_reader", -] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py index ea331aa3ce..0a6765bd39 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py @@ -5,7 +5,9 @@ Each reader is a stateless handler for one on-disk format, safe to reuse across files. Readers are looked up by ``file_format`` (from :func:`detect_format`) so the pipeline never branches on format -itself. +itself. Built-in readers self-register (a ``register_reader`` call at the bottom of each module) and +are loaded lazily on the first :func:`get_reader` call, so importing this module does not pull in +pyarrow until a reader is actually needed. """ from __future__ import annotations @@ -15,7 +17,6 @@ from typing import Any, ClassVar, Protocol import pyarrow as pa - from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource @@ -40,6 +41,20 @@ def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = No _READERS: dict[str, FormatReader] = {} +_builtins_loaded = False + + +def _load_builtin_readers() -> None: + """Import the built-in reader modules so their self-registration runs (once). + + Deferred to call time — not import time — so there is no cycle with the reader modules that import + from this one, and pyarrow stays out of the import graph until a reader is actually resolved. + """ + global _builtins_loaded + if _builtins_loaded: + return + _builtins_loaded = True + from nemo_datasets_plugin.profiler.readers import jsonl, parquet # noqa: F401 self-registering def register_reader(reader: FormatReader) -> None: @@ -47,6 +62,7 @@ def register_reader(reader: FormatReader) -> None: def get_reader(file_format: str) -> FormatReader: + _load_builtin_readers() try: return _READERS[file_format] except KeyError: diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py index 607f231860..f59b420c37 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py @@ -21,7 +21,10 @@ def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = No stripped = raw_line.strip() if not stripped: # tolerate blank lines between records continue - rows.append(json.loads(stripped)) + record = json.loads(stripped) + if not isinstance(record, dict): + continue # a record is a column map; skip stray scalars/arrays rather than crash downstream + rows.append(record) if row_cap is not None and len(rows) >= row_cap: break diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py index b360e18e40..bea95ce212 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py @@ -6,7 +6,6 @@ from __future__ import annotations import pyarrow.parquet as pq - from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource from nemo_datasets_plugin.profiler.readers.base import ReadResult, register_reader diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index b072dce32f..e2eec1f86f 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -189,7 +189,9 @@ def _content_len(content: Any) -> int: if isinstance(content, str): return len(content) if isinstance(content, list): # VLM content as a list of typed parts - return sum(len(part.get("text", "")) for part in content if isinstance(part, dict)) + return sum( + len(part["text"]) for part in content if isinstance(part, dict) and isinstance(part.get("text"), str) + ) return 0 diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index 235be47945..13990c3a16 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -153,3 +153,27 @@ def test_implicit_prompt_evidence_from_embedded_transcript(): result = classify(features, {}, rows) assert result.prompt_form == "implicit" assert any(e.kind == "content_probe" for e in result.evidence) + + +def test_ground_truth_may_be_a_container_dtype(): + # test_cases (list) and verification_info (struct) are verification targets, not free text, + # so the text-only dtype gate must not drop them. + features = [_f("prompt", "string"), _f("test_cases", "list"), _f("verification_info", "struct")] + classify(features, {}) + assert features[1].semantic_role == "ground_truth" + assert features[2].semantic_role == "ground_truth" + + +def test_container_ground_truth_drives_verifiability(): + features = [_f("prompt", "string"), _f("test_cases", "list")] + rows = [{"prompt": "q", "test_cases": [{"in": "1", "out": "2"}]}, {"prompt": "q2", "test_cases": []}] + result = classify(features, {}, rows) + assert result.verifiability.method == "ground_truth_column" + assert result.verifiability.coverage == 0.5 # the empty test_cases list is not a usable target + + +def test_bare_scalar_ground_truth_alias_is_still_rejected(): + # A numeric column named "ground_truth" is far more likely a label/score than a target. + features = [_f("ground_truth", "int64")] + classify(features, {}) + assert features[0].semantic_role is None diff --git a/plugins/nemo-datasets/tests/test_cli.py b/plugins/nemo-datasets/tests/test_cli.py index 65e2d3886e..deb335486a 100644 --- a/plugins/nemo-datasets/tests/test_cli.py +++ b/plugins/nemo-datasets/tests/test_cli.py @@ -3,12 +3,24 @@ """Smoke tests for the ``nemo datasets`` CLI surface.""" +import json + +import pyarrow as pa +import pyarrow.parquet as pq +import typer from nemo_datasets_plugin.cli import DatasetsCLI from typer.testing import CliRunner runner = CliRunner() +def _mounted() -> typer.Typer: + """The app as the platform mounts it: ``nemo datasets ``.""" + root = typer.Typer() + root.add_typer(DatasetsCLI().get_cli(), name="datasets") + return root + + def test_cli_metadata(): cli = DatasetsCLI() assert cli.name == "datasets" @@ -20,3 +32,28 @@ def test_profile_command_registered(): result = runner.invoke(app, ["--help"]) assert result.exit_code == 0 assert "profile" in result.stdout + + +def test_profile_command_profiles_a_directory(tmp_path): + pq.write_table( + pa.Table.from_pylist([{"prompt": "a"}, {"prompt": "b"}]), + tmp_path / "train-00000-of-00001.parquet", + ) + result = runner.invoke(_mounted(), ["datasets", "profile", str(tmp_path)]) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + assert payload["partitions"][0]["features"][0]["name"] == "prompt" + assert payload["sampling"]["rows_scanned"] == 2 + + +def test_profile_command_rejects_unknown_output_format(tmp_path): + result = runner.invoke(_mounted(), ["datasets", "profile", str(tmp_path), "--output", "xml"]) + assert result.exit_code != 0 + + +def test_profile_command_rejects_non_directory(tmp_path): + target = tmp_path / "not-a-dir" + target.write_text("x") + result = runner.invoke(_mounted(), ["datasets", "profile", str(target)]) + assert result.exit_code != 0 diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 2cded58b87..7ecda36763 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -7,10 +7,10 @@ import pyarrow as pa import pyarrow.parquet as pq -from nemo_datasets_plugin.profiler import profile from nemo_datasets_plugin.profiler.digest import content_digest from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource from nemo_datasets_plugin.profiler.partition import group_partitions +from nemo_datasets_plugin.profiler.pipeline import profile from nemo_datasets_plugin.profiler.splits import resolve_splits FIXED_TIME = datetime(2026, 7, 13, 12, 0, 0, tzinfo=timezone.utc) @@ -180,3 +180,48 @@ def test_profile_is_deterministic(tmp_path): first = profile(source, created_at=FIXED_TIME) second = profile(source, created_at=FIXED_TIME) assert first.model_dump_json() == second.model_dump_json() + + +def test_profile_tolerates_non_object_jsonl_lines(tmp_path): + # A valid-JSON-but-non-object line parses cleanly, so the reader (not the read) must handle it; + # otherwise it would poison the unprotected schema/stats stage and abort the whole profile. + (tmp_path / "train.jsonl").write_text('{"a": 1}\n[1, 2, 3]\n{"a": 2}\n') + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.partitions[0].splits[0].num_examples == 2 # objects counted, stray array dropped + assert result.sampling.exhaustive is True + + +def test_profile_digest_covers_only_stored_files(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + without_card = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + (tmp_path / "README.md").write_text("a dataset card") # a non-data file, never stored as a FileRecord + with_card = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + # A file the profile does not store must not move the digest... + assert without_card.content_digest == with_card.content_digest + # ...and the digest is recomputable from exactly the FileRecords the profile stores. + stored = [ + FileEntry(path=f.path, size_bytes=f.size_bytes, checksum=f.checksum) + for partition in with_card.partitions + for split in partition.splits + for f in split.files + ] + assert content_digest(stored) == with_card.content_digest + + +def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): + # If detect_format recognizes an extension the registry has no reader for, that file must be + # isolated like a corrupt one, not crash the whole profile. + from nemo_datasets_plugin.profiler.readers import base + + monkeypatch.setitem(base._EXTENSION_FORMATS, ".xyz", "xyz-no-reader") + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + (tmp_path / "extra.xyz").write_text("whatever") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise + + records = {f.path: f for p in result.partitions for s in p.splits for f in s.files} + assert records["extra.xyz"].num_rows is None # kept, but unreadable + assert result.sampling.exhaustive is False diff --git a/plugins/nemo-datasets/tests/test_readers.py b/plugins/nemo-datasets/tests/test_readers.py index 0c7ffed45b..08973522d2 100644 --- a/plugins/nemo-datasets/tests/test_readers.py +++ b/plugins/nemo-datasets/tests/test_readers.py @@ -7,7 +7,7 @@ import pyarrow.parquet as pq import pytest from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource -from nemo_datasets_plugin.profiler.readers import detect_format, get_reader +from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader PARQUET_ROWS = [ {"prompt": "a", "score": 1}, @@ -115,3 +115,12 @@ def test_jsonl_reader_row_cap_leaves_count_unknown(tmp_path): assert result.rows == [{"a": 1}, {"a": 2}] assert result.rows_scanned == 2 assert result.num_rows is None # a partial read can't assert the total + + +def test_jsonl_reader_skips_non_object_lines(tmp_path): + # A record is a column map; valid JSON that is a scalar or array is not a row. + (tmp_path / "d.jsonl").write_text('{"a": 1}\n[1, 2, 3]\n42\n"loose"\n{"a": 2}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)) + + assert result.rows == [{"a": 1}, {"a": 2}] # stray non-object lines dropped, objects kept + assert result.num_rows == 2 diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 07a1a2cfba..7735831971 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -87,6 +87,13 @@ def test_message_ends_with_user_turn_is_prompt_only_signal(): assert stats.messages.ends_with_assistant_rate == 0.0 +def test_message_content_parts_tolerate_non_string_text(): + # A VLM-style content part whose "text" key is present but not a string must not crash measurement. + rows = [{"m": [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": None}]}]}] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.content_chars.max == 0 # no measurable text, and no crash + + # --- sparsity and null rate ---------------------------------------------------------------------- From d99e2a97f57182549c6c2bf5f34e4d718afb2b00 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Tue, 14 Jul 2026 00:19:30 -0400 Subject: [PATCH 09/44] fix(datasets): gate verifiability on a minimum coverage Assert a verification method only when its coverage clears a floor (_MIN_VERIFIABILITY_COVERAGE = 0.05), and fall through between methods so a sparse ground_truth column no longer masks a strong extractable-answer signal. Previously a single coincidental "#### " match (e.g. 1/1500 rows in trl-lib/kto-mix-14k) asserted extractable_final_answer at ~0% coverage. Add tests for the below-floor drop, above-floor assertion, and the ground_truth -> extractable fall-through. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/classify.py | 22 +++++++++++---- plugins/nemo-datasets/tests/test_classify.py | 28 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index ef29e2f036..2ebe79503d 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -195,6 +195,11 @@ def has(*required: str) -> bool: _GSM8K_ANSWER = re.compile(r"####\s*-?[\d.,/]+\s*$") _BOXED_ANSWER = re.compile(r"\\boxed\{") +# A verification target must cover at least this fraction of sampled rows to be asserted. Below it, +# a "hit" is noise -- e.g. one completion in thousands coincidentally ending in `#### ` does +# not make a dataset verifiable. Tune here; the coverage itself is still reported on the Verifiability. +_MIN_VERIFIABILITY_COVERAGE = 0.05 + def _pct(fraction: float) -> str: return f"{round(fraction * 100)}%" @@ -220,20 +225,25 @@ def _detect_verifiability(features: list[FeatureSchema], rows: list[dict]) -> Ve if not rows: return None + # Each method wins only if it clears the coverage floor; otherwise fall through to the next, so a + # sparse ground_truth column can still yield to an extractable-answer signal instead of masking it. ground_truth = next((feature for feature in features if feature.semantic_role == "ground_truth"), None) if ground_truth is not None: present = sum(1 for row in rows if row.get(ground_truth.name) not in (None, "", [])) coverage = present / len(rows) - detail = f"'{ground_truth.name}' present in {_pct(coverage)} of {len(rows)} sampled rows" - return Verifiability( - method="ground_truth_column", coverage=coverage, evidence=[Evidence(kind="content_probe", detail=detail)] - ) + if coverage >= _MIN_VERIFIABILITY_COVERAGE: + detail = f"'{ground_truth.name}' present in {_pct(coverage)} of {len(rows)} sampled rows" + return Verifiability( + method="ground_truth_column", + coverage=coverage, + evidence=[Evidence(kind="content_probe", detail=detail)], + ) texts = _completion_texts(features, rows) if texts: hits = sum(1 for text in texts if _GSM8K_ANSWER.search(text) or _BOXED_ANSWER.search(text)) - if hits: - coverage = hits / len(texts) + coverage = hits / len(texts) + if coverage >= _MIN_VERIFIABILITY_COVERAGE: detail = f"completion ends with an extractable answer (#### or \\boxed) in {_pct(coverage)} of {len(texts)} sampled rows" return Verifiability( method="extractable_final_answer", diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index 13990c3a16..47d8f5062a 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -147,6 +147,34 @@ def test_no_verifiability_without_a_target(): assert result.verifiability is None +def test_verifiability_ignores_below_threshold_extractable_noise(): + # One coincidental "#### " in a large sample is noise, not a verifiable dataset (kto-mix-14k). + features = [_f("prompt", "string"), _f("completion", "string")] + rows = [{"prompt": "q", "completion": "just prose"} for _ in range(100)] + rows[0]["completion"] = "the answer is #### 7" # 1/100 = 1% < 5% floor + assert classify(features, {}, rows).verifiability is None + + +def test_verifiability_asserted_above_coverage_floor(): + features = [_f("prompt", "string"), _f("completion", "string")] + rows = [{"prompt": "q", "completion": "just prose"} for _ in range(10)] + for row in rows[:2]: + row["completion"] = "answer #### 7" # 2/10 = 20% >= 5% floor + result = classify(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 0.2 + + +def test_sparse_ground_truth_falls_through_to_extractable_answer(): + # A ground_truth column present in too few rows must not mask a strong extractable-answer signal. + features = [_f("completion", "string"), _f("ground_truth", "string")] + rows = [{"completion": "reasoning #### 5", "ground_truth": None} for _ in range(100)] + rows[0]["ground_truth"] = "5" # 1/100 ground_truth coverage -> below floor, must fall through + result = classify(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + + def test_implicit_prompt_evidence_from_embedded_transcript(): features = [_f("chosen", "string"), _f("rejected", "string")] rows = [{"chosen": "\n\nHuman: hi\n\nAssistant: hello", "rejected": "\n\nHuman: hi\n\nAssistant: hey"}] From 3ff7bb02ff3c103afd75b74de6b36157ddb7a6df Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Tue, 14 Jul 2026 11:27:50 -0400 Subject: [PATCH 10/44] Make `nemo datasets profile` a real command Signed-off-by: Albert Cui --- plugins/nemo-datasets/pyproject.toml | 2 +- plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py | 7 +++++++ pyproject.toml | 7 +++++-- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/plugins/nemo-datasets/pyproject.toml b/plugins/nemo-datasets/pyproject.toml index 89075e4bc0..4a803821fd 100644 --- a/plugins/nemo-datasets/pyproject.toml +++ b/plugins/nemo-datasets/pyproject.toml @@ -12,7 +12,7 @@ dependencies = [ version = "0.1.0" [tool.uv.sources] -nemo-platform-plugin = { path = "../../packages/nemo_platform_plugin", editable = true } +nemo-platform-plugin = { workspace = true } [project.entry-points."nemo.cli"] datasets = "nemo_datasets_plugin.cli:DatasetsCLI" diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py index 497b7eecbb..75fbb24613 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py @@ -18,6 +18,13 @@ class DatasetsCLI(NemoCLI): def get_cli(self) -> typer.Typer: app = typer.Typer(help="Dataset profiling commands.") + @app.callback() + def _root() -> None: + """Dataset profiling commands.""" + # A no-op callback keeps ``profile`` an explicit subcommand (``nemo datasets profile + # ``) instead of Typer collapsing the lone command into ``nemo datasets ``, + # which would break the moment a second command is added. + @app.command() def profile( path: str = typer.Argument(..., help="Path to a local directory of dataset files."), diff --git a/pyproject.toml b/pyproject.toml index a89e3a8e54..5217d7b606 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -184,6 +184,7 @@ enabled-plugins = [ "nemo-switchyard", "nemo-agents-plugin", "nemo-deployments-plugin[docker,k8s]", + "nemo-datasets-plugin", "nemo-customizer-plugin", "nemo-automodel-plugin", "nemo-optimization-plugin", @@ -405,6 +406,7 @@ nemo-safe-synthesizer-plugin = { workspace = true } nemo-switchyard = { workspace = true } nemo-agents-plugin = { workspace = true } nemo-deployments-plugin = { workspace = true } +nemo-datasets-plugin = { workspace = true } nemo-agents-example-calculator = { workspace = true } nemo-agents-example-email-phishing = { workspace = true } nemo-agents-example-email-security = { workspace = true } @@ -463,6 +465,7 @@ members = [ "plugins/nemo-safe-synthesizer", "plugins/nemo-switchyard", "plugins/nemo-agents", + "plugins/nemo-datasets", "plugins/nemo-deployments", "plugins/nemo-insights", "plugins/nemo-eval-author", @@ -585,8 +588,8 @@ extra-paths = [ # example-plugin is not a workspace member; add its src so ty can # resolve nmp.example_plugin imports when checking plugin test files. "plugins/example-plugin/src", - # nemo-datasets plugin is not a workspace member yet; add its src so ty can - # resolve nemo_datasets_plugin imports when checking plugin test files. + # nemo-datasets plugin is a workspace member; add its src so ty can + # resolve nemo_datasets_plugin imports from plugin test files. "plugins/nemo-datasets/src", # nemo-guardrails plugin is not a workspace member; add its src so ty can # resolve nemo_guardrails_plugin imports when checking plugin test files. From 1110c44cab126a174489c66b2d3fc2502e71fa1f Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 16 Jul 2026 15:36:10 -0400 Subject: [PATCH 11/44] resync uv.lock Signed-off-by: Albert Cui --- uv.lock | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/uv.lock b/uv.lock index cf41b73710..a80424c030 100644 --- a/uv.lock +++ b/uv.lock @@ -31,6 +31,7 @@ members = [ "nemo-automodel-plugin", "nemo-customizer-plugin", "nemo-data-designer-plugin", + "nemo-datasets-plugin", "nemo-deployments-plugin", "nemo-eval-author-plugin", "nemo-evaluator-plugin", @@ -4397,6 +4398,27 @@ requires-dist = [ ] provides-extras = ["test", "data-designer-nemo", "nemo-platform-plugin"] +[[package]] +name = "nemo-datasets-plugin" +version = "0.1.0" +source = { editable = "plugins/nemo-datasets" } +dependencies = [ + { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyarrow", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pyyaml", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, + { name = "pyarrow", specifier = ">=17.0.0" }, + { name = "pydantic", specifier = ">=2.10.3" }, + { name = "pyyaml", specifier = ">=6.0.2" }, + { name = "typer", specifier = ">=0.20.0,<0.26" }, +] + [[package]] name = "nemo-deployments-plugin" version = "0.0.0" @@ -6641,6 +6663,7 @@ core-services = [ { name = "nemo-automodel-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-customizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-data-designer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-datasets-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-deployments-plugin", extra = ["docker", "k8s"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-eval-author-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-evaluator-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6740,6 +6763,7 @@ enabled-plugins = [ { name = "nemo-automodel-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-customizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-data-designer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-datasets-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-deployments-plugin", extra = ["docker", "k8s"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-eval-author-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-evaluator-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6763,6 +6787,7 @@ functional-services = [ { name = "nemo-automodel-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-customizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-data-designer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-datasets-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-deployments-plugin", extra = ["docker", "k8s"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-eval-author-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-evaluator-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6861,6 +6886,7 @@ core-services = [ { name = "nemo-automodel-plugin", editable = "plugins/nemo-automodel" }, { name = "nemo-customizer-plugin", editable = "plugins/nemo-customizer" }, { name = "nemo-data-designer-plugin", editable = "plugins/nemo-data-designer" }, + { name = "nemo-datasets-plugin", editable = "plugins/nemo-datasets" }, { name = "nemo-deployments-plugin", extras = ["docker", "k8s"], editable = "plugins/nemo-deployments" }, { name = "nemo-eval-author-plugin", marker = "python_full_version < '3.14'", editable = "plugins/nemo-eval-author" }, { name = "nemo-evaluator-plugin", editable = "plugins/nemo-evaluator" }, @@ -6963,6 +6989,7 @@ enabled-plugins = [ { name = "nemo-automodel-plugin", editable = "plugins/nemo-automodel" }, { name = "nemo-customizer-plugin", editable = "plugins/nemo-customizer" }, { name = "nemo-data-designer-plugin", editable = "plugins/nemo-data-designer" }, + { name = "nemo-datasets-plugin", editable = "plugins/nemo-datasets" }, { name = "nemo-deployments-plugin", extras = ["docker", "k8s"], editable = "plugins/nemo-deployments" }, { name = "nemo-eval-author-plugin", marker = "python_full_version < '3.14'", editable = "plugins/nemo-eval-author" }, { name = "nemo-evaluator-plugin", editable = "plugins/nemo-evaluator" }, @@ -6986,6 +7013,7 @@ functional-services = [ { name = "nemo-automodel-plugin", editable = "plugins/nemo-automodel" }, { name = "nemo-customizer-plugin", editable = "plugins/nemo-customizer" }, { name = "nemo-data-designer-plugin", editable = "plugins/nemo-data-designer" }, + { name = "nemo-datasets-plugin", editable = "plugins/nemo-datasets" }, { name = "nemo-deployments-plugin", extras = ["docker", "k8s"], editable = "plugins/nemo-deployments" }, { name = "nemo-eval-author-plugin", marker = "python_full_version < '3.14'", editable = "plugins/nemo-eval-author" }, { name = "nemo-evaluator-plugin", editable = "plugins/nemo-evaluator" }, From 0ada117a6d817989bf750a4a49626ea105758dbb Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 24 Jul 2026 16:54:26 -0400 Subject: [PATCH 12/44] fix(datasets): drop non-finite floats from numeric column stats A NaN or +-inf in a numeric column poisoned NumericStats.min/max/mean. Pydantic serializes a non-finite float to JSON null, which then fails to re-validate against the required float fields -- so the whole profile became unreadable on the next load (e.g. GET .../profile). Filter to finite values before computing the numeric summary; an all-non-finite column yields no numeric block rather than a poisoned one. Signed-off-by: Albert Cui --- .../src/nemo_datasets_plugin/profiler/stats.py | 9 ++++++++- plugins/nemo-datasets/tests/test_stats.py | 16 +++++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index e2eec1f86f..13acf5d7de 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -57,7 +57,14 @@ def _column_stats(feature: FeatureSchema, values: list[Any], total: int, exhaust if counts is not None and counts.distinct_count <= _MAX_ENUM_VALUES: categorical = counts # a bounded string enumeration, not free text elif _is_numeric(feature.dtype): - numbers = [float(value) for value in present if isinstance(value, (int, float)) and not isinstance(value, bool)] + # Drop non-finite floats (NaN / +-inf): they serialize to JSON null and then fail to + # re-validate against NumericStats' required floats, which would make the whole profile + # unreadable on the next load. + numbers = [ + float(value) + for value in present + if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) + ] if numbers: numeric = NumericStats(min=min(numbers), max=max(numbers), mean=sum(numbers) / len(numbers)) categorical = _cardinality(present, exhaustive) diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 7735831971..af23406d7c 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -4,7 +4,7 @@ """Tests for per-column statistics.""" from nemo_datasets_plugin.profiler.stats import derive_stats -from nemo_platform_plugin.files.dataset_profile import FeatureSchema +from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema def _feature(name, dtype): @@ -58,6 +58,20 @@ def test_numeric_cardinality_values_withheld_when_not_exhaustive(): assert stats.categorical.values is None # a sample cannot prove the enumeration +def test_numeric_stats_ignore_non_finite_values(): + # NaN / +-inf poison min/max/mean and serialize to JSON null, which then fails to re-validate + # against NumericStats' required floats -- making the whole profile unreadable. Drop them. + values = [1.0, float("nan"), 3.0, float("inf"), float("-inf"), 5.0] + stats = derive_stats([_feature("n", "float64")], _rows("n", values), exhaustive=True)["n"] + assert (stats.numeric.min, stats.numeric.max, stats.numeric.mean) == (1.0, 5.0, 3.0) + ColumnStats.model_validate_json(stats.model_dump_json()) # round-trips: no NaN/inf leaked into JSON + + +def test_numeric_all_non_finite_yields_no_numeric_summary(): + stats = derive_stats([_feature("n", "float64")], _rows("n", [float("nan"), float("inf")]), exhaustive=True) + assert stats.get("n") is None or stats["n"].numeric is None + + # --- messages ------------------------------------------------------------------------------------ From 58a2b1a926c185fb3d2fea5939c78b07a3464ef2 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 24 Jul 2026 16:54:32 -0400 Subject: [PATCH 13/44] fix(datasets): profile mixed-format directories as separate partitions group_partitions groups by directory only, so a directory holding more than one format (a stray .jsonl beside .parquet shards) was profiled as a single partition -- deriving features from whichever format was read first while measuring rows from both. Sub-group each directory's files by format so every partition is format-homogeneous; a directory that actually splits gets format-qualified partition names (e.g. default:parquet, default:jsonl) while single-format directories keep their bare name. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/pipeline.py | 161 +++++++++++------- plugins/nemo-datasets/tests/test_pipeline.py | 20 +++ 2 files changed, 118 insertions(+), 63 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 683ee9e70d..744861125d 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -19,7 +19,7 @@ from nemo_datasets_plugin.profiler.classify import classify from nemo_datasets_plugin.profiler.digest import content_digest -from nemo_datasets_plugin.profiler.file_source import FileSource +from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader from nemo_datasets_plugin.profiler.schema import derive_features @@ -61,69 +61,17 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase all_scanned = True for partition_name, partition_entries in group_partitions(data_entries): - partition_rows: list[dict] = [] - arrow_schema = None - partition_scanned = True - split_profiles: list[SplitProfile] = [] - for split in resolve_splits(partition_entries): - file_records: list[FileRecord] = [] - split_examples = 0 - split_counts_known = True # every file's exact total row count is known (footer or full scan) - split_scanned = True # every row of every file was actually parsed - for entry in split.entries: - try: - result = get_reader(_format_of(entry.path)).read(source, entry) - except Exception: - # Failure isolation: an unreadable file (or missing reader) keeps its identity, - # skips its rows, and does not abort the profile. - result = None - if result is None: - num_rows = None - scanned_all = False - else: - num_rows = result.num_rows - rows_scanned += result.rows_scanned - partition_rows.extend(result.rows) - if arrow_schema is None: - arrow_schema = result.arrow_schema - # Exhaustive requires parsing every row; a known footer count alone is not enough. - scanned_all = num_rows is not None and result.rows_scanned >= num_rows - file_records.append( - FileRecord( - path=entry.path, - size_bytes=entry.size_bytes, - checksum=entry.checksum, - num_rows=num_rows, - ) - ) - if num_rows is None: - split_counts_known = False - else: - split_examples += num_rows - if not scanned_all: - split_scanned = False - all_scanned = all_scanned and split_scanned - partition_scanned = partition_scanned and split_scanned - split_profiles.append( - SplitProfile( - name=split.name, - canonical=split.canonical, - files=file_records, - num_examples=split_examples if split_counts_known else None, - ) - ) - features = derive_features(partition_rows, arrow_schema) - stats = derive_stats(features, partition_rows, exhaustive=partition_scanned) - partitions.append( - PartitionProfile( - name=partition_name, - file_format=_format_of(partition_entries[0].path), - splits=split_profiles, - features=features, - stats=stats, - classification=classify(features, stats, partition_rows), + format_groups = _split_by_format(partition_entries) + for file_format, format_entries in format_groups: + # A directory that holds more than one format yields one partition per format; qualify + # the name so the partitions stay distinct. A single-format directory keeps its bare name. + name = f"{partition_name}:{file_format}" if len(format_groups) > 1 else partition_name + partition, partition_rows_scanned, partition_scanned = _profile_partition( + source, name, file_format, format_entries ) - ) + partitions.append(partition) + rows_scanned += partition_rows_scanned + all_scanned = all_scanned and partition_scanned sampling = SamplingInfo( exhaustive=all_scanned, @@ -142,3 +90,90 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase sampling=sampling, partitions=partitions, ) + + +def _split_by_format(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry]]]: + """Sub-group a directory's files by format so each profiled partition is format-homogeneous. + + ``group_partitions`` groups by directory only, but one directory can hold more than one format + (a stray ``.jsonl`` beside ``.parquet`` shards). Left mixed, a partition would derive its schema + from whichever format was read first and then measure rows from both. Sorted for deterministic + partition order; ``entries`` are pre-filtered ``data_entries`` so every format is registered. + """ + by_format: dict[str, list[FileEntry]] = {} + for entry in entries: + by_format.setdefault(_format_of(entry.path), []).append(entry) + return sorted(by_format.items()) + + +def _profile_partition( + source: FileSource, name: str, file_format: str, entries: list[FileEntry] +) -> tuple[PartitionProfile, int, bool]: + """Profile one format-homogeneous partition. + + Returns the partition plus its ``(rows_scanned, scanned_all)`` contribution to the dataset-level + sampling envelope. An unreadable file (or a format with no registered reader) is isolated: it + keeps its FileRecord, contributes no rows, and flips ``scanned_all`` off — it never aborts. + """ + partition_rows: list[dict] = [] + arrow_schema = None + rows_scanned = 0 + partition_scanned = True + split_profiles: list[SplitProfile] = [] + for split in resolve_splits(entries): + file_records: list[FileRecord] = [] + split_examples = 0 + split_counts_known = True # every file's exact total row count is known (footer or full scan) + split_scanned = True # every row of every file was actually parsed + for entry in split.entries: + try: + result = get_reader(file_format).read(source, entry) + except Exception: + # Failure isolation: an unreadable file (or missing reader) keeps its identity, + # skips its rows, and does not abort the profile. + result = None + if result is None: + num_rows = None + scanned_all = False + else: + num_rows = result.num_rows + rows_scanned += result.rows_scanned + partition_rows.extend(result.rows) + if arrow_schema is None: + arrow_schema = result.arrow_schema + # Exhaustive requires parsing every row; a known footer count alone is not enough. + scanned_all = num_rows is not None and result.rows_scanned >= num_rows + file_records.append( + FileRecord( + path=entry.path, + size_bytes=entry.size_bytes, + checksum=entry.checksum, + num_rows=num_rows, + ) + ) + if num_rows is None: + split_counts_known = False + else: + split_examples += num_rows + if not scanned_all: + split_scanned = False + partition_scanned = partition_scanned and split_scanned + split_profiles.append( + SplitProfile( + name=split.name, + canonical=split.canonical, + files=file_records, + num_examples=split_examples if split_counts_known else None, + ) + ) + features = derive_features(partition_rows, arrow_schema) + stats = derive_stats(features, partition_rows, exhaustive=partition_scanned) + partition = PartitionProfile( + name=name, + file_format=file_format, + splits=split_profiles, + features=features, + stats=stats, + classification=classify(features, stats, partition_rows), + ) + return partition, rows_scanned, partition_scanned diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 7ecda36763..6c9fc93208 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -147,6 +147,26 @@ def test_profile_multiple_directories_become_partitions(tmp_path): assert all(p.file_format == "parquet" for p in result.partitions) +def test_profile_splits_mixed_formats_into_separate_partitions(tmp_path): + # A directory holding two formats must not profile as one partition: features/stats would be + # derived from one format's schema but measured over rows from both. Each format becomes its own + # format-homogeneous partition, name-qualified so the two stay distinct. + _write_parquet(tmp_path / "data" / "train-00000-of-00001.parquet", [{"prompt": "a"}]) + (tmp_path / "data" / "extra.jsonl").write_text('{"question": "b"}\n{"question": "c"}\n') + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + by_format = {p.file_format: p for p in result.partitions} + assert set(by_format) == {"parquet", "jsonl"} + assert by_format["parquet"].name == "default:parquet" + assert by_format["jsonl"].name == "default:jsonl" + # Each partition's schema reflects only its own files. + assert [f.name for f in by_format["parquet"].features] == ["prompt"] + assert [f.name for f in by_format["jsonl"].features] == ["question"] + assert result.sampling.rows_scanned == 3 # 1 parquet + 2 jsonl, each counted once + assert result.sampling.exhaustive is True + + def test_profile_isolates_unreadable_files(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) (tmp_path / "test-00000-of-00001.parquet").write_bytes(b"not a real parquet file") From 8faf743953851a18aa5ab016e2adc9702a28eb96 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 31 Jul 2026 17:58:47 -0400 Subject: [PATCH 14/44] feat(datasets): record per-file read errors in the profile contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A profile could say a file had no row count, but never why. That left "corrupt input", "format we cannot read yet", and "profiler bug" looking identical to a consumer, which is the worst case for the one artifact meant to save people from re-inspecting a dataset by hand. Add `FileRecord.error`, and an `error` value to the `Evidence.kind` vocabulary for when a detector could not run at all — so an absent finding stays distinguishable from a finding of absence. Also narrow `PartitionProfile.file_format`: it advertised csv and arrow as though they were read today. They are reserved vocabulary; the profiler reports files in those formats as unsupported rather than profiling them. Both additions are optional and backward compatible, so PROFILE_SCHEMA_VERSION is left at 1.0. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index d2038ea280..3ea4e69f20 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -46,7 +46,11 @@ class Evidence(BaseModel): """ kind: str = Field( - description="column_name | column_dtype | content_probe | split_name | file_name | card_metadata", + description=( + "column_name | column_dtype | content_probe | split_name | file_name | card_metadata | " + "error — the last for when a detector could not run at all, so an absent finding is " + "distinguishable from a finding of absence." + ), ) detail: str = Field( description="Self-describing evidence, e.g. \"answer matches '#### ' in 100% of 1024 sampled rows\".", @@ -262,6 +266,14 @@ class FileRecord(BaseModel): default=None, description="Exact only (parquet footer / exhaustive scan), else None.", ) + error: str | None = Field( + default=None, + description=( + "Why this file was not fully read, when it wasn't — unreadable, corrupt, or partially " + "parsed. None means a clean read. Without it a missing `num_rows` is indistinguishable " + "from a profiler bug, and a consumer cannot tell corrupt input from unsupported input." + ), + ) class SplitProfile(BaseModel): @@ -315,7 +327,13 @@ class PartitionProfile(BaseModel): """ name: str = "default" - file_format: str = Field(description="jsonl | parquet | csv | arrow") + file_format: str = Field( + description=( + "jsonl | parquet are read today; csv | arrow are reserved vocabulary the profiler cannot " + "read yet — files in those formats are reported as unsupported rather than profiled, so " + "this value never appears without a real partition behind it." + ), + ) splits: list[SplitProfile] = Field(description="card-declared > path-detected > single 'default' split.") features: list[FeatureSchema] = Field( description="The row schema: measured layout plus detected role markers, derived de novo (nested).", From d045d49ca881cbc919e1225fa9a08f1f24070a98 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 31 Jul 2026 17:59:01 -0400 Subject: [PATCH 15/44] fix(datasets): keep reading a JSONL file past a corrupt line MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One truncated line used to cost the entire file. json.loads raised, the pipeline's blanket `except Exception` dropped the whole read, and the file lost its row count and every column it was the only witness for — silently, since nothing recorded a reason. Skip the unparseable line instead, count it, and report the first failure on the new ReadResult.error channel so the pipeline can mark the file as only partly understood. A line of valid JSON that simply is not a row (a stray scalar or array) is deliberately not an error: it is not a row of this dataset, so the count stays exact and the scan stays exhaustive. Only data we failed to parse counts against us. Also add `is_unsupported_data`, naming the extensions that plainly hold records but have no reader yet, so the pipeline can distinguish "data I cannot read" from "not data" rather than dropping both on the floor. Signed-off-by: Albert Cui --- .../profiler/readers/base.py | 15 ++++++++++ .../profiler/readers/jsonl.py | 29 ++++++++++++++++--- plugins/nemo-datasets/tests/test_readers.py | 21 ++++++++++++++ 3 files changed, 61 insertions(+), 4 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py index 0a6765bd39..4afb575a80 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py @@ -28,6 +28,10 @@ class ReadResult: rows_scanned: int # number of rows actually parsed num_rows: int | None = None # exact total when cheaply known (e.g. a parquet footer), else None arrow_schema: pa.Schema | None = None # the declared column schema, when the format carries one + # Why the read understood less than the whole file, when that happened. None means nothing was + # lost. This is the only channel a reader has to explain a partial result, so a consumer can tell + # "corrupt input" from "unsupported format" from "profiler bug" instead of seeing a silent gap. + error: str | None = None class FormatReader(Protocol): @@ -79,3 +83,14 @@ def get_reader(file_format: str) -> FormatReader: def detect_format(path: str) -> str | None: """Map a file path to a registered format by extension, or None when unrecognized.""" return _EXTENSION_FORMATS.get(Path(path).suffix.lower()) + + +# Extensions that plainly hold dataset records but have no reader yet. Naming them explicitly is what +# lets the profiler say "there is data here I cannot read" instead of treating a dataset it does not +# understand as an empty one — a README or a LICENSE is genuinely not data and stays ignored. +_UNSUPPORTED_DATA_EXTENSIONS = {".csv", ".tsv", ".arrow", ".feather", ".json", ".avro", ".orc"} + + +def is_unsupported_data(path: str) -> bool: + """Whether a path looks like dataset records this profiler has no reader for.""" + return Path(path).suffix.lower() in _UNSUPPORTED_DATA_EXTENSIONS diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py index f59b420c37..85a12395a0 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py @@ -16,20 +16,41 @@ class JsonlReader: def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: rows: list[dict] = [] + unparseable = 0 + first_failure: str | None = None + hit_cap = False with source.open(entry.path) as stream: - for raw_line in stream: + for line_number, raw_line in enumerate(stream, start=1): stripped = raw_line.strip() if not stripped: # tolerate blank lines between records continue - record = json.loads(stripped) + try: + record = json.loads(stripped) + except ValueError as exc: + # A truncated or corrupt line costs that line, never the file. Dropping the whole + # file would erase its row count and every column it was the only witness for. + unparseable += 1 + if first_failure is None: + first_failure = f"line {line_number}: {exc}" + continue if not isinstance(record, dict): continue # a record is a column map; skip stray scalars/arrays rather than crash downstream rows.append(record) if row_cap is not None and len(rows) >= row_cap: + hit_cap = True break - num_rows = len(rows) if row_cap is None else None - return ReadResult(rows=rows, rows_scanned=len(rows), num_rows=num_rows, arrow_schema=None) + # `num_rows` counts records the reader could parse. A line of valid JSON that simply is not a + # row (a stray scalar or array) is not a row of this dataset, so it leaves the count exact and + # sets no error. An *unparseable* line is data we failed to read, so it is reported: the + # pipeline reads `error` to decide the file was not exhaustively scanned. + # + # A cap only costs the exact count when it actually stopped the read. A file smaller than the + # cap was still read to EOF, so it keeps an exact count — which is what lets a capped profile + # of a small dataset stay exhaustive instead of degrading every stat for no reason. + num_rows = None if hit_cap else len(rows) + error = f"skipped {unparseable} unparseable line(s); first at {first_failure}" if unparseable else None + return ReadResult(rows=rows, rows_scanned=len(rows), num_rows=num_rows, arrow_schema=None, error=error) register_reader(JsonlReader()) diff --git a/plugins/nemo-datasets/tests/test_readers.py b/plugins/nemo-datasets/tests/test_readers.py index 08973522d2..391b11aa5a 100644 --- a/plugins/nemo-datasets/tests/test_readers.py +++ b/plugins/nemo-datasets/tests/test_readers.py @@ -75,6 +75,7 @@ def test_parquet_reader_reads_schema_rows_and_exact_count(tmp_path): assert result.num_rows == 3 # exact, from the footer assert result.rows_scanned == 3 assert result.rows == PARQUET_ROWS + assert result.arrow_schema is not None assert set(result.arrow_schema.names) == {"prompt", "score"} @@ -124,3 +125,23 @@ def test_jsonl_reader_skips_non_object_lines(tmp_path): assert result.rows == [{"a": 1}, {"a": 2}] # stray non-object lines dropped, objects kept assert result.num_rows == 2 + # Not a read failure: those lines are not rows of this dataset, so the count stays exact and the + # file is still exhaustively scanned. Only an *unparseable* line is an error. + assert result.error is None + + +def test_jsonl_reader_survives_an_unparseable_line(tmp_path): + # One truncated line must cost that line, not the file. Dropping the whole file would erase its + # row count and any column it was the only witness for. + (tmp_path / "d.jsonl").write_text('{"a": 1}\n{"a": 2\n{"a": 3}\n') + result = get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)) + + assert result.rows == [{"a": 1}, {"a": 3}] # the readable rows survive + assert result.rows_scanned == 2 + assert result.error is not None + assert "line 2" in result.error # self-describing: which line, and why + + +def test_jsonl_reader_clean_read_reports_no_error(tmp_path): + (tmp_path / "d.jsonl").write_text('{"a": 1}\n{"a": 2}\n') + assert get_reader("jsonl").read(LocalFileSource(tmp_path), FileEntry("d.jsonl", 0)).error is None From 96a4470f18034f647e4623d382fd9286e5e173fd Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 31 Jul 2026 17:59:20 -0400 Subject: [PATCH 16/44] fix(datasets): measure chat rows and column cardinality correctly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five measurement bugs, all of which made the profile quietly wrong rather than loudly broken. ShareGPT data was invisible. Only `{role, content}` counted as a chat message, so `{from, value}` stayed a plain `list`, failed the `messages` dtype gate in classification, and profiled as `unknown` with no stats at all — over a large slice of public chat data. Any chat not using the literal role "assistant" looked prompt-only. Matching one string meant gpt/bot/model conversations reported ends_with_assistant_rate 0.0, which classification reads as "no training target". Match an assistant-equivalent set instead. roles_seen still reports what was actually there, verbatim: the contract is explicit that an unexpected role is the finding worth surfacing, not something to normalize. has_tool_calls was true for any parquet chat schema that merely declared the field. to_pylist materializes every declared struct field as an explicit None, so `"tool_calls" in message` matched on every row of datasets that never make a tool call. Test the value, not the key. A non-string role aborted the entire profile. roles_seen is typed list[str], so a numeric role raised a ValidationError from inside the one stage nothing guarded. Coerce it. distinct_count was withheld for strings above 32 distinct values — dropping the id-like signal (~= row count) precisely where it carries the most information — and bool columns got no stats at all, leaving the very column that decides unpaired_preference with no measured class balance. Finally, _text_quality ran three interpreted passes over every character and dominated profiling time. Scanning with the regex engine instead measures the same three signals identically (verified across empty, unicode, run-heavy and alternating inputs) at 32 MB/s rather than 12.6 MB/s. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/schema.py | 18 ++-- .../nemo_datasets_plugin/profiler/stats.py | 96 ++++++++++++++----- plugins/nemo-datasets/tests/test_schema.py | 14 +++ plugins/nemo-datasets/tests/test_stats.py | 45 ++++++++- 4 files changed, 142 insertions(+), 31 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py index 393798c4b3..92d4c2c508 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py @@ -5,9 +5,9 @@ Derive the ``features`` tree (a list of :class:`FeatureSchema`) de novo from the data. Parquet carries a declared schema, so it is converted directly; formats without one (jsonl) are inferred -from the sampled rows by resolving each column's dtype. A list of ``{role, content}`` structs is -recognized as the ``messages`` dtype, and a list whose elements are all the same length records that -length as ``fixed_length``. +from the sampled rows by resolving each column's dtype. A list of ``{role, content}`` structs — or +ShareGPT's ``{from, value}`` spelling of the same thing — is recognized as the ``messages`` dtype, +and a list whose elements are all the same length records that length as ``fixed_length``. """ from __future__ import annotations @@ -17,8 +17,11 @@ import pyarrow as pa from nemo_platform_plugin.files.dataset_profile import FeatureSchema -# A list element carrying at least these keys is treated as a chat message (the messages dtype). -_MESSAGE_KEYS = {"role", "content"} +# A list element carrying at least one of these key sets is treated as a chat message (the messages +# dtype). ShareGPT-style data spells the same structure `{from, value}`; recognizing only +# `{role, content}` left a large slice of public chat data typed as a plain `list`, which then failed +# the `messages` dtype gate in classification and profiled as `unknown` with no stats at all. +_MESSAGE_KEY_SETS = ({"role", "content"}, {"from", "value"}) def derive_features(rows: list[dict[str, Any]], arrow_schema: pa.Schema | None = None) -> list[FeatureSchema]: @@ -32,7 +35,10 @@ def derive_features(rows: list[dict[str, Any]], arrow_schema: pa.Schema | None = def _is_message_struct(item: FeatureSchema) -> bool: - return item.dtype == "struct" and item.fields is not None and _MESSAGE_KEYS <= {field.name for field in item.fields} + if item.dtype != "struct" or item.fields is None: + return False + names = {field.name for field in item.fields} + return any(keys <= names for keys in _MESSAGE_KEY_SETS) # --- from a declared arrow schema (parquet) ------------------------------------------------------ diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 13acf5d7de..1b97eea3e6 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -13,6 +13,7 @@ from __future__ import annotations import math +import re from typing import Any from nemo_platform_plugin.files.dataset_profile import ( @@ -37,6 +38,10 @@ def derive_stats( total = len(rows) stats: dict[str, ColumnStats] = {} for feature in features: + # Parquet permits duplicate field names, and stats is keyed by name. Measuring the first and + # skipping the rest makes which one wins deterministic instead of "whichever came last". + if feature.name in stats: + continue column = _column_stats(feature, [row.get(feature.name) for row in rows], total, exhaustive) if column is not None: stats[feature.name] = column @@ -53,9 +58,15 @@ def _column_stats(feature: FeatureSchema, values: list[Any], total: int, exhaust if strings: text = TextStats(chars=_quantiles([len(value) for value in strings])) quality = _text_quality(strings) - counts = _cardinality(present, exhaustive) - if counts is not None and counts.distinct_count <= _MAX_ENUM_VALUES: - categorical = counts # a bounded string enumeration, not free text + # distinct_count is always safe to store and is the id-like signal (~= rows_scanned) the + # contract documents; only the values themselves are row data, and _cardinality already gates + # those on an exhaustive read. Withholding the count for high-cardinality strings dropped the + # signal precisely where it carries the most information. + categorical = _cardinality(present, exhaustive) + elif feature.dtype == "bool": + # The column that decides unpaired_preference deserves a measured class balance rather than + # no stats at all. + categorical = _cardinality(present, exhaustive) elif _is_numeric(feature.dtype): # Drop non-finite floats (NaN / +-inf): they serialize to JSON null and then fail to # re-validate against NumericStats' required floats, which would make the whole profile @@ -111,6 +122,14 @@ def _cardinality(present: list[Any], exhaustive: bool) -> CategoricalStats | Non # --- text quality -------------------------------------------------------------------------------- +_WHITESPACE_RUN = re.compile(r"\s") +_NON_ASCII_RUN = re.compile(r"[^\x00-\x7f]") +# Any character repeated four or more times in a row. Scanning with the regex engine instead of a +# Python loop is what keeps this affordable: these three measurements used to run three interpreted +# passes over every character of every string and dominated total profiling time. +_REPEAT_RUN = re.compile(r"(.)\1{3,}", re.DOTALL) + + def _text_quality(strings: list[str]) -> TextQuality: total_chars = 0 whitespace = 0 @@ -118,8 +137,9 @@ def _text_quality(strings: list[str]) -> TextQuality: repetition_sum = 0.0 for value in strings: total_chars += len(value) - whitespace += sum(char.isspace() for char in value) - non_ascii += sum(ord(char) > 127 for char in value) + # str.count-style scanning in C rather than a per-character generator in Python. + whitespace += _count_matches(_WHITESPACE_RUN, value) + non_ascii += _count_matches(_NON_ASCII_RUN, value) repetition_sum += _repetition_score(value) return TextQuality( whitespace_ratio=whitespace / total_chars if total_chars else 0.0, @@ -128,6 +148,10 @@ def _text_quality(strings: list[str]) -> TextQuality: ) +def _count_matches(pattern: re.Pattern[str], text: str) -> int: + return sum(1 for _ in pattern.finditer(text)) + + def _repetition_score(text: str) -> float: """Fraction of characters inside a run of the same character of length >= 4. @@ -136,22 +160,40 @@ def _repetition_score(text: str) -> float: """ if not text: return 0.0 - redundant = 0 - run = 1 - for index in range(1, len(text)): - if text[index] == text[index - 1]: - run += 1 - else: - if run >= 4: - redundant += run - run = 1 - if run >= 4: - redundant += run + redundant = sum(len(match.group(0)) for match in _REPEAT_RUN.finditer(text)) return redundant / len(text) # --- messages ------------------------------------------------------------------------------------ +# Role strings that mean "the turn the model is trained to produce". Matching only the literal +# "assistant" made every chat dataset using another convention (ShareGPT's gpt, or bot/model) look +# like it ended on a user turn, which classification reads as a prompt-only dataset with no training +# target — a false negative over a large slice of public chat data. +_ASSISTANT_ROLES = {"assistant", "gpt", "bot", "model", "chatbot", "ai"} + + +def _message_field(message: dict, *names: str) -> Any: + """The first present, non-null value among ``names``. + + Chat rows spell the same two fields either ``{role, content}`` or ``{from, value}``. Reading with + a plain ``.get`` default is not enough: parquet materializes *every* declared struct field, so an + absent field arrives as an explicit None rather than a missing key. + """ + for name in names: + value = message.get(name) + if value is not None: + return value + return None + + +def _role_of(message: dict) -> Any: + return _message_field(message, "role", "from") + + +def _is_assistant_role(role: Any) -> bool: + return isinstance(role, str) and role.lower() in _ASSISTANT_ROLES + def _message_stats(rows_messages: list[list]) -> MessageStats | None: if not rows_messages: @@ -169,14 +211,22 @@ def _message_stats(rows_messages: list[list]) -> MessageStats | None: for message in messages: if not isinstance(message, dict): continue - role = message.get("role") - if role is not None and role not in roles_seen: - roles_seen.append(role) - total_content += _content_len(message.get("content")) - if "tool_calls" in message or role == "tool": + role = _role_of(message) + if role is not None: + # Coerced to str because roles_seen is typed list[str] and a non-string role would + # fail validation — aborting the whole profile from inside the one stage the pipeline + # does not guard. Reported verbatim otherwise: the contract is explicit that an + # unexpected role is the finding worth surfacing, not something to normalize away. + role = role if isinstance(role, str) else str(role) + if role not in roles_seen: + roles_seen.append(role) + total_content += _content_len(_message_field(message, "content", "value")) + # `.get` truthiness, not `in`: parquet materializes every declared struct field, so a + # schema that merely declares tool_calls would otherwise report tool use on every row. + if message.get("tool_calls") or role == "tool": has_tool_calls = True content_chars.append(total_content) - if messages and isinstance(messages[-1], dict) and messages[-1].get("role") == "assistant": + if messages and isinstance(messages[-1], dict) and _is_assistant_role(_role_of(messages[-1])): ends_with_assistant += 1 if _valid_alternation(messages): valid_alternation += 1 @@ -204,5 +254,5 @@ def _content_len(content: Any) -> int: def _valid_alternation(messages: list) -> bool: """True when user/assistant turns alternate (ignoring any leading system turns).""" - roles = [m.get("role") for m in messages if isinstance(m, dict) and m.get("role") != "system"] + roles = [_role_of(m) for m in messages if isinstance(m, dict) and _role_of(m) != "system"] return all(roles[i] != roles[i + 1] for i in range(len(roles) - 1)) diff --git a/plugins/nemo-datasets/tests/test_schema.py b/plugins/nemo-datasets/tests/test_schema.py index 71f3868b8a..e492ffcfe2 100644 --- a/plugins/nemo-datasets/tests/test_schema.py +++ b/plugins/nemo-datasets/tests/test_schema.py @@ -56,6 +56,20 @@ def test_from_rows_list_of_role_content_structs_is_messages(): assert {f.name for f in feature.items.fields} == {"role", "content"} +def test_from_rows_sharegpt_from_value_is_messages(): + # ShareGPT spells the same structure {from, value}. Recognizing only {role, content} left it a + # plain list, which then failed the messages dtype gate and profiled as `unknown` with no stats. + rows = [{"conversations": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "yo"}]}] + feature = derive_features(rows)[0] + assert feature.dtype == "messages" + assert {f.name for f in feature.items.fields} == {"from", "value"} + + +def test_from_arrow_sharegpt_from_value_is_messages(): + schema = pa.schema([("conversations", pa.list_(pa.struct([("from", pa.string()), ("value", pa.string())])))]) + assert derive_features([], schema)[0].dtype == "messages" + + def test_from_rows_constant_length_list_records_fixed_length(): feature = derive_features([{"e": [0.1, 0.2, 0.3]}, {"e": [0.4, 0.5, 0.6]}])[0] assert feature.dtype == "list" diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index af23406d7c..4599cce5db 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -33,15 +33,24 @@ def test_text_quality_flags_repetition_and_non_ascii(): assert stats.quality.non_ascii_ratio > 0.0 # accented characters -def test_free_text_string_has_no_categorical_but_low_cardinality_does(): +def test_string_cardinality_counts_always_but_withholds_free_text_values(): + # distinct_count is the id-like signal and is always safe to store; only the values themselves + # are row data, and only a small proven enumeration may be kept. free_text = derive_stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)]), exhaustive=True) - assert free_text["t"].categorical is None # too many distinct values to be an enumeration + assert free_text["t"].categorical.distinct_count == 50 # ~= row count -> id-like + assert free_text["t"].categorical.values is None # too many distinct values to be an enumeration labels = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"]), exhaustive=True) assert labels["c"].categorical.distinct_count == 2 assert labels["c"].categorical.values == ["no", "yes"] # proven enumeration under exhaustive read +def test_bool_column_gets_a_measured_class_balance(): + stats = derive_stats([_feature("label", "bool")], _rows("label", [True, False, True]), exhaustive=True) + assert stats["label"].categorical.distinct_count == 2 + assert stats["label"].categorical.values == ["False", "True"] + + # --- numeric ------------------------------------------------------------------------------------- @@ -101,6 +110,38 @@ def test_message_ends_with_user_turn_is_prompt_only_signal(): assert stats.messages.ends_with_assistant_rate == 0.0 +def test_message_stats_read_sharegpt_from_value(): + rows = [{"m": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello there"}]}] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.roles_seen == ["human", "gpt"] # verbatim, not normalized + assert stats.messages.content_chars.max == len("hi") + len("hello there") + assert stats.messages.ends_with_assistant_rate == 1.0 # "gpt" is the responder turn + + +def test_assistant_equivalent_roles_count_as_the_training_target(): + # Matching only the literal "assistant" made every other convention look prompt-only. + for responder in ("assistant", "gpt", "bot", "model", "AI"): + rows = [{"m": [{"role": "user", "content": "q"}, {"role": responder, "content": "a"}]}] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.ends_with_assistant_rate == 1.0, responder + + +def test_non_string_role_does_not_break_measurement(): + # roles_seen is typed list[str]; a numeric role used to raise a ValidationError from inside the + # one stage the pipeline did not guard, aborting the whole profile. + rows = [{"m": [{"role": 1, "content": "hi"}]}] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.roles_seen == ["1"] + + +def test_declared_but_unset_tool_calls_is_not_tool_use(): + # parquet materializes every declared struct field, so `"tool_calls" in message` reported tool + # use for any schema that merely declares the field. + rows = [{"m": [{"role": "user", "content": "hi", "tool_calls": None}]}] + stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + assert stats.messages.has_tool_calls is False + + def test_message_content_parts_tolerate_non_string_text(): # A VLM-style content part whose "text" key is present but not a string must not crash measurement. rows = [{"m": [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": None}]}]}] From 8a0f8cc170de307fd889b33bfb70e06ecf03e5ae Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 31 Jul 2026 17:59:35 -0400 Subject: [PATCH 17/44] fix(datasets): require rank to have something to rank, accept integer labels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `rank` was checked second in the type ladder, ahead of every more specific structure, so any numeric column named rank decided the dataset type on its own. chosen+rejected+rank came out ranked_responses instead of a preference pair; so did prompt+response+score+rank; so did a lone rank column with nothing to rank at all. Require a completion-ish role alongside it, and order it below the structures it was overriding. A `label` column only counted when boolean, which made unpaired_preference effectively unreachable — 0/1 integers are how these datasets actually ship. Accept an integer label when the observed values really are binary. A wider integer range is a class index or a rating, which is a different claim, so it stays unroled; the cardinality check is possible now that bool and integer columns carry distinct_count. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/classify.py | 42 ++++++++++----- plugins/nemo-datasets/tests/test_classify.py | 51 ++++++++++++++++++- 2 files changed, 79 insertions(+), 14 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index 2ebe79503d..915cf33e92 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -88,13 +88,30 @@ def _is_numeric(dtype: str) -> bool: return dtype.startswith(("int", "uint", "float")) -def _role_for(feature: FeatureSchema) -> str | None: +def _is_binary(column: ColumnStats | None) -> bool: + """Whether a column was observed to hold at most two distinct values.""" + return column is not None and column.categorical is not None and column.categorical.distinct_count <= 2 + + +def _is_label_column(feature: FeatureSchema, stats: dict[str, ColumnStats]) -> bool: + """Whether a column named ``label`` really carries a binary preference label. + + A bool says so outright. An integer is the more common on-disk encoding (0/1, as KTO-style sets + ship it), but only when the observed values really are binary — a wider integer range is a class + index or a rating, which is a different claim, so it stays unroled. + """ + if feature.dtype == "bool": + return True + return _is_numeric(feature.dtype) and _is_binary(stats.get(feature.name)) + + +def _role_for(feature: FeatureSchema, stats: dict[str, ColumnStats]) -> str | None: name = feature.name.lower() dtype = feature.dtype if name in _SCORE_ALIASES and _is_numeric(dtype): return "score" - if name == "label" and dtype == "bool": - return "label" + if name == "label": + return "label" if _is_label_column(feature, stats) else None role = _ALIAS_ROLES.get(name) if role is None: @@ -111,14 +128,12 @@ def _role_for(feature: FeatureSchema) -> str | None: return None if role in {"stepwise_completions", "stepwise_labels"} and dtype != "list": return None - if role == "label": # "label" reached here only when not bool - return None return role -def _assign_roles(features: list[FeatureSchema]) -> None: +def _assign_roles(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> None: for feature in features: - role = _role_for(feature) + role = _role_for(feature, stats) if role is not None: feature.semantic_role = role @@ -167,14 +182,17 @@ def has(*required: str) -> bool: if has("prompt", "stepwise_completions", "stepwise_labels"): return "stepwise_supervision" - if has("rank"): - return "ranked_responses" + if has("chosen", "rejected"): + return "preference_pair" if has("prompt", "completion", "score"): return "scored_response" if has("prompt", "completion", "label"): return "unpaired_preference" - if has("chosen", "rejected"): - return "preference_pair" + # `rank` is only a dataset type alongside something to rank. On its own it short-circuited every + # more specific structure above, so a stray numeric column named `rank` — or a ranked variant of + # a preference set — was enough to mislabel the dataset. + if has("rank") and roles & {"completion", "chosen", "rejected", "stepwise_completions"}: + return "ranked_responses" if has("prompt", "completion"): return "prompt_completion" if "messages" in roles: @@ -299,7 +317,7 @@ def classify( detection; role/axis/type inference needs only the schema and stats. """ rows = rows or [] - _assign_roles(features) + _assign_roles(features, stats) roles = {feature.semantic_role for feature in features if feature.semantic_role} dataset_type = _detect_type(features, stats) fmt = _detect_format(features) diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index 47d8f5062a..6035e632c8 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -4,7 +4,13 @@ """Tests for classification: role assignment, format/prompt-form axes, and dataset type.""" from nemo_datasets_plugin.profiler.classify import classify -from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema, MessageStats, Quantiles +from nemo_platform_plugin.files.dataset_profile import ( + CategoricalStats, + ColumnStats, + FeatureSchema, + MessageStats, + Quantiles, +) def _f(name, dtype): @@ -81,11 +87,52 @@ def test_scored_response_beats_prompt_completion(): assert classify(features, {}).dataset_type == "scored_response" -def test_unpaired_preference_needs_boolean_label(): +def test_unpaired_preference_accepts_a_boolean_label(): features = [_f("prompt", "string"), _f("completion", "string"), _f("label", "bool")] assert classify(features, {}).dataset_type == "unpaired_preference" +def test_unpaired_preference_accepts_a_binary_integer_label(): + # 0/1 is the usual on-disk encoding; requiring a bool made unpaired_preference unreachable for + # most real datasets. + features = [_f("prompt", "string"), _f("completion", "string"), _f("label", "int64")] + stats = {"label": ColumnStats(categorical=CategoricalStats(distinct_count=2))} + assert classify(features, stats).dataset_type == "unpaired_preference" + assert features[2].semantic_role == "label" + + +def test_wide_integer_label_is_not_a_preference_label(): + # A multi-class index or a rating is a different claim from a binary preference. + features = [_f("prompt", "string"), _f("completion", "string"), _f("label", "int64")] + stats = {"label": ColumnStats(categorical=CategoricalStats(distinct_count=7))} + assert classify(features, stats).dataset_type == "prompt_completion" + assert features[2].semantic_role is None + + +# --- rank ------------------------------------------------------------------------------------ + + +def test_rank_needs_something_to_rank(): + # A lone numeric column named "rank" used to short-circuit every more specific structure. + features = [_f("rank", "int64")] + assert classify(features, {}).dataset_type == "unknown" + + +def test_rank_does_not_override_a_preference_pair(): + features = [_f("chosen", "string"), _f("rejected", "string"), _f("rank", "int64")] + assert classify(features, {}).dataset_type == "preference_pair" + + +def test_rank_does_not_override_scored_responses(): + features = [_f("prompt", "string"), _f("response", "string"), _f("helpfulness", "int64"), _f("rank", "int64")] + assert classify(features, {}).dataset_type == "scored_response" + + +def test_rank_alongside_a_completion_is_ranked_responses(): + features = [_f("prompt", "string"), _f("completion", "string"), _f("rank", "int64")] + assert classify(features, {}).dataset_type == "ranked_responses" + + def test_messages_ending_on_assistant_is_messages_type(): result = classify([_f("messages", "messages")], {"messages": _messages_column(1.0)}) assert result.dataset_type == "messages" From 0a305e83f7aa13971fb988c792729e7ecf6ea815 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 31 Jul 2026 17:59:59 -0400 Subject: [PATCH 18/44] fix(datasets): correct partition, split and sampling structure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The structural half of the profile had four problems that made it either wrong or dishonest. Schema was taken from whichever shard sorted first. A column appearing only in a later shard vanished from features, and so from stats. Worse, the same data classified differently depending on file order: prompt_completion one way round, scored_response the other. Unify the partition's schemas instead, which is order-independent; on a genuine type conflict fall back to inferring from the rows, widening the disputed column to json rather than asserting one shard's type over the other's. Splits ignored directories entirely. Both common HuggingFace layouts were mis-modelled: top-level train/ and test/ became two unrelated partitions, each with its own schema and classification, while data/train/ and data/test/ collapsed into a single "default" split with train and test rows pooled into one stats blob. Read a split-named directory off the path, and stop treating one as a partition dimension. The envelope claimed things that were not true. A directory of .csv shards profiled as exhaustive with rows_total 0 — indistinguishable from a dataset that really is empty — files_scanned counted files that were never opened, and per-file read failures had nowhere to be reported. Report unsupported files, count only successful reads, propagate FileRecord.error, and leave rows_total unknown rather than lying with a 0 the contract forbids anyway. Finally, nothing bounded memory. Reads materialized every row of every file as Python dicts at ~6.5x the on-disk parquet size, so `nemo datasets profile` on a real dataset was an OOM: 1 GB in, ~6 GB resident. Read every file, but cap rows per file (sampling a subset of *files* would hide columns that only appear in later shards). Files under the cap are still read to EOF and keep exact counts, so small datasets stay exhaustive and lose nothing. On a 21.6 MB fixture this is 0.06s and 38 MB against 2.85s and 141 MB, with the exact row count still read from the footer. `strategy` records the policy while `exhaustive` records the outcome, which is why the contract keeps them separate. Schema derivation, stats and classification also now run behind a guard: they are pure computation over rows already in memory, and leaving them unguarded meant one odd value could abort an otherwise complete profile from the only stage nothing was catching. A failure costs that partition its measurements, says so as `error` evidence, and keeps its structure. Signed-off-by: Albert Cui --- .../src/nemo_datasets_plugin/cli.py | 15 +- .../profiler/partition.py | 13 +- .../nemo_datasets_plugin/profiler/pipeline.py | 173 +++++++++--- .../nemo_datasets_plugin/profiler/splits.py | 47 +++- plugins/nemo-datasets/tests/test_pipeline.py | 258 +++++++++++++++++- 5 files changed, 455 insertions(+), 51 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py index 75fbb24613..7adf17aa49 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py @@ -29,8 +29,18 @@ def _root() -> None: def profile( path: str = typer.Argument(..., help="Path to a local directory of dataset files."), output: str = typer.Option("json", "--output", "-o", help="Output format: json | yaml."), + rows_per_file: int = typer.Option( + None, + "--rows-per-file", + help="Rows to read from each file (default 1000); 0 reads every row, which is exact " + "but scales memory with the dataset rather than the file count.", + min=0, + ), ) -> None: """Profile a local dataset directory and print its DatasetProfile.""" + # Imported here, not at module scope: the platform calls get_cli() for every plugin at + # startup, and the profiler pulls in pyarrow. The row-cap default lives in the pipeline + # rather than being restated here, so an unspecified flag simply omits the argument. from nemo_datasets_plugin.profiler.file_source import LocalFileSource from nemo_datasets_plugin.profiler.pipeline import profile as run_profile @@ -41,7 +51,10 @@ def profile( except NotADirectoryError as exc: raise typer.BadParameter(str(exc)) from exc - result = run_profile(source) + if rows_per_file is None: + result = run_profile(source) + else: + result = run_profile(source, row_cap=rows_per_file or None) if output == "yaml": import yaml diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py index ab2b9557c9..33c7c84b97 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py @@ -12,12 +12,21 @@ from pathlib import PurePosixPath from nemo_datasets_plugin.profiler.file_source import FileEntry +from nemo_datasets_plugin.profiler.splits import is_split_directory def _top_dir(path: str) -> str | None: - """The first path segment when the file is nested, else None for a root-level file.""" + """The partition directory for a file: its first path segment, else None for a root-level file. + + A split-named top-level directory (``train/``, ``test/``) is deliberately *not* a partition + dimension. Grouping on it would split one dataset's train and test into unrelated partitions, + each deriving its own schema and classification — the exact structure `splits` exists to model. + Those files fall through to the same partition and are separated by :mod:`splits` instead. + """ parts = PurePosixPath(path).parts - return parts[0] if len(parts) > 1 else None + if len(parts) <= 1 or is_split_directory(parts[0]): + return None + return parts[0] def group_partitions(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry]]]: diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 744861125d..61c2a0d1a0 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -9,25 +9,33 @@ (``features``), per-column ``stats``, and the full ``classification`` (roles, format, prompt form, dataset type, and verifiability). -Reads are exhaustive (every row of every file). Sampling large datasets with bounded probes is a -later, drop-in optimization behind the same reader seam. +Every file is opened — sampling a subset of files would hide columns that appear only in later shards +— but each is read up to ``row_cap`` rows, so peak memory tracks the file count rather than the +dataset size. Files smaller than the cap are read to the end and keep their exact counts, so capping +costs nothing on a small dataset. Pass ``row_cap=None`` for a genuinely exhaustive scan. """ from __future__ import annotations +from dataclasses import dataclass from datetime import datetime, timezone +import pyarrow as pa from nemo_datasets_plugin.profiler.classify import classify from nemo_datasets_plugin.profiler.digest import content_digest from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource from nemo_datasets_plugin.profiler.partition import group_partitions -from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader +from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import resolve_splits from nemo_datasets_plugin.profiler.stats import derive_stats from nemo_platform_plugin.files.dataset_profile import ( + ColumnStats, DatasetProfile, + Evidence, + FeatureSchema, FileRecord, + PartitionClassification, PartitionProfile, SamplingInfo, SplitProfile, @@ -36,6 +44,14 @@ PROFILER_NAME = "nemo-dataset-profiler" PROFILER_VERSION = "0.1.0" +# Rows read per file by default. Every file is still opened — head-sampling a *subset of files* would +# hide columns that appear only in later shards — but each is capped, so peak memory scales with the +# file count rather than the dataset size. Uncapped, a partition materializes every row of every file +# as Python dicts at roughly 6x the on-disk parquet size, which puts a 10 GB dataset far past any +# reasonable machine. A thousand rows per file is ample for the statistics computed here (length +# quantiles, rates, cardinality); pass ``row_cap=None`` for a genuinely exhaustive scan. +DEFAULT_ROW_CAP = 1000 + def _format_of(path: str) -> str: """The registered format of a data file. Callers pass only pre-filtered ``data_entries``, so the @@ -46,18 +62,34 @@ def _format_of(path: str) -> str: return file_format -def profile(source: FileSource, *, created_at: datetime | None = None) -> DatasetProfile: +def profile( + source: FileSource, + *, + created_at: datetime | None = None, + row_cap: int | None = DEFAULT_ROW_CAP, +) -> DatasetProfile: """Profile the dataset behind ``source`` into a ``DatasetProfile``. + ``row_cap`` bounds how many rows are read from each file; ``None`` reads every row, which is + exact but scales memory with the dataset rather than the file count. Files smaller than the cap + are still read to the end, so a capped profile of a small dataset stays exhaustive. + ``created_at`` is injectable so a profile can be made reproducible byte-for-byte in tests; it defaults to the current UTC time. """ created_at = created_at or datetime.now(timezone.utc) all_entries = source.list_files() data_entries = [entry for entry in all_entries if detect_format(entry.path) is not None] + # Files that plainly hold records but have no reader yet. They are not profiled, but they must be + # reported: silently dropping them let a directory of .csv shards profile as an exhaustively + # scanned, empty dataset — indistinguishable from a dataset that really is empty. + unsupported = sorted( + entry.path for entry in all_entries if detect_format(entry.path) is None and is_unsupported_data(entry.path) + ) partitions: list[PartitionProfile] = [] rows_scanned = 0 + files_read = 0 all_scanned = True for partition_name, partition_entries in group_partitions(data_entries): @@ -66,27 +98,38 @@ def profile(source: FileSource, *, created_at: datetime | None = None) -> Datase # A directory that holds more than one format yields one partition per format; qualify # the name so the partitions stay distinct. A single-format directory keeps its bare name. name = f"{partition_name}:{file_format}" if len(format_groups) > 1 else partition_name - partition, partition_rows_scanned, partition_scanned = _profile_partition( - source, name, file_format, format_entries - ) - partitions.append(partition) - rows_scanned += partition_rows_scanned - all_scanned = all_scanned and partition_scanned + outcome = _profile_partition(source, name, file_format, format_entries, row_cap) + partitions.append(outcome.partition) + rows_scanned += outcome.rows_scanned + files_read += outcome.files_read + all_scanned = all_scanned and outcome.scanned_all + + # Data we could not read is data we did not scan, so unsupported files defeat exhaustiveness just + # as an unreadable file does. + exhaustive = all_scanned and not unsupported + profiler_info: dict = {"name": PROFILER_NAME, "version": PROFILER_VERSION} + if unsupported: + profiler_info["unsupported_files"] = unsupported sampling = SamplingInfo( - exhaustive=all_scanned, - strategy="full", + exhaustive=exhaustive, + # The policy in effect, which `exhaustive` deliberately does not encode: a capped run over + # files that all fit under the cap is still a full scan, and an uncapped run can still fall + # short of exhaustive because a file was unreadable. + strategy="full" if row_cap is None else "head_per_file", + # `rows_total` is documented as never zero: a 0 here would read as "this dataset is empty" + # when it more often means nothing was recognized. Unknown is the honest answer. + rows_total=rows_scanned if exhaustive and rows_scanned else None, rows_scanned=rows_scanned, - rows_total=rows_scanned if all_scanned else None, - files_scanned=len(data_entries), - per_file_row_cap=None, - seed=None, + files_scanned=files_read, # files actually opened and read, not files merely listed + per_file_row_cap=row_cap, + seed=None, # head sampling makes no random choices; a seed would be theatre ) return DatasetProfile( # Digest only the files stored as FileRecords, so the profile can recompute its own digest. content_digest=content_digest(data_entries), created_at=created_at, - profiler_info={"name": PROFILER_NAME, "version": PROFILER_VERSION}, + profiler_info=profiler_info, sampling=sampling, partitions=partitions, ) @@ -106,18 +149,72 @@ def _split_by_format(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry return sorted(by_format.items()) +def _unify_schemas(schemas: list[pa.Schema]) -> pa.Schema | None: + """One schema describing every file of the partition, or None when they cannot be reconciled. + + Taking the first file's schema and ignoring the rest makes the profile depend on which shard + happens to sort first: a column that appears only in a later shard would vanish from ``features`` + (and so from ``stats``), and the same data would classify differently under a different file + order. Unifying is order-independent for the common case — later shards adding columns. + + A genuine type conflict for the same column name has no correct answer here, so we return None + and let the caller fall back to inferring from the rows themselves, which widens the conflicting + column to ``json`` rather than asserting one shard's type over the other's. + """ + if not schemas: + return None + if len(schemas) == 1: + return schemas[0] + try: + return pa.unify_schemas(schemas) + except pa.ArrowException: + return None + + +def _measure( + partition_rows: list[dict], arrow_schemas: list[pa.Schema], exhaustive: bool +) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: + """Derive schema, stats and classification, degrading to structure-only if any of it fails. + + These three stages are pure computation over rows already in memory, so a failure here is either + a profiler bug or data shaped in a way no detector anticipated. Reads are already isolated per + file; leaving this stage unguarded meant one odd value — a chat message whose ``role`` is a number + — could abort an otherwise complete profile from the one place nothing was catching. The + partition's structure (files, splits, row counts) is established by then and stays useful, so the + failure costs its measurements and says so, rather than the entire run. + """ + try: + features = derive_features(partition_rows, _unify_schemas(arrow_schemas)) + stats = derive_stats(features, partition_rows, exhaustive=exhaustive) + return features, stats, classify(features, stats, partition_rows) + except Exception as exc: + detail = f"could not measure this partition: {type(exc).__name__}: {exc}" + return [], {}, PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]) + + +@dataclass(frozen=True) +class _PartitionOutcome: + """One partition plus what it contributes to the dataset-level sampling envelope.""" + + partition: PartitionProfile + rows_scanned: int + files_read: int # files actually opened and read, so `files_scanned` can exclude failures + scanned_all: bool + + def _profile_partition( - source: FileSource, name: str, file_format: str, entries: list[FileEntry] -) -> tuple[PartitionProfile, int, bool]: + source: FileSource, name: str, file_format: str, entries: list[FileEntry], row_cap: int | None +) -> _PartitionOutcome: """Profile one format-homogeneous partition. - Returns the partition plus its ``(rows_scanned, scanned_all)`` contribution to the dataset-level - sampling envelope. An unreadable file (or a format with no registered reader) is isolated: it - keeps its FileRecord, contributes no rows, and flips ``scanned_all`` off — it never aborts. + An unreadable file (or a format with no registered reader) is isolated: it keeps its FileRecord, + records *why* on ``FileRecord.error``, contributes no rows, and flips ``scanned_all`` off — it + never aborts the profile. """ partition_rows: list[dict] = [] - arrow_schema = None + arrow_schemas: list[pa.Schema] = [] rows_scanned = 0 + files_read = 0 partition_scanned = True split_profiles: list[SplitProfile] = [] for split in resolve_splits(entries): @@ -126,29 +223,36 @@ def _profile_partition( split_counts_known = True # every file's exact total row count is known (footer or full scan) split_scanned = True # every row of every file was actually parsed for entry in split.entries: + error: str | None = None try: - result = get_reader(file_format).read(source, entry) - except Exception: + result = get_reader(file_format).read(source, entry, row_cap=row_cap) + except Exception as exc: # Failure isolation: an unreadable file (or missing reader) keeps its identity, - # skips its rows, and does not abort the profile. + # skips its rows, and does not abort the profile. The reason is recorded rather than + # swallowed, so a consumer can tell corrupt input from a profiler bug. result = None + error = f"{type(exc).__name__}: {exc}" if result is None: num_rows = None scanned_all = False else: + files_read += 1 + error = result.error num_rows = result.num_rows rows_scanned += result.rows_scanned partition_rows.extend(result.rows) - if arrow_schema is None: - arrow_schema = result.arrow_schema - # Exhaustive requires parsing every row; a known footer count alone is not enough. - scanned_all = num_rows is not None and result.rows_scanned >= num_rows + if result.arrow_schema is not None: + arrow_schemas.append(result.arrow_schema) + # Exhaustive requires parsing every row; a known footer count alone is not enough, and + # a partial read (corrupt lines skipped) is not exhaustive however many rows it got. + scanned_all = num_rows is not None and result.rows_scanned >= num_rows and error is None file_records.append( FileRecord( path=entry.path, size_bytes=entry.size_bytes, checksum=entry.checksum, num_rows=num_rows, + error=error, ) ) if num_rows is None: @@ -166,14 +270,15 @@ def _profile_partition( num_examples=split_examples if split_counts_known else None, ) ) - features = derive_features(partition_rows, arrow_schema) - stats = derive_stats(features, partition_rows, exhaustive=partition_scanned) + features, stats, classification = _measure(partition_rows, arrow_schemas, partition_scanned) partition = PartitionProfile( name=name, file_format=file_format, splits=split_profiles, features=features, stats=stats, - classification=classify(features, stats, partition_rows), + classification=classification, + ) + return _PartitionOutcome( + partition=partition, rows_scanned=rows_scanned, files_read=files_read, scanned_all=partition_scanned ) - return partition, rows_scanned, partition_scanned diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py index ae582a0f20..6b5c1afc9a 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py @@ -4,8 +4,9 @@ """Split resolution from file paths. Given the files in one partition, group them into splits by inferring each file's split from its -name. A declared split map from a dataset card would take precedence over this inference, but card -parsing is not wired up yet, so path inference is the only source today. +path — a split-named directory first, then the shard-stripped filename. A declared split map from a +dataset card would take precedence over this inference, but card parsing is not wired up yet, so +path inference is the only source today. """ from __future__ import annotations @@ -16,8 +17,10 @@ from nemo_datasets_plugin.profiler.file_source import FileEntry -# Strips a shard suffix like "-00000" or "-00000-of-00003" from a file stem. -_SHARD_SUFFIX = re.compile(r"-\d{2,}(?:-of-\d{2,})?$") +# Strips a shard suffix like "-00000" or "-00000-of-00003" from a file stem. A bare trailing number +# must be zero-padded to count: `-\d{2,}` alone also matched years and versions, turning +# covid-19.jsonl into a "covid" split and data-2024.jsonl into "data". +_SHARD_SUFFIX = re.compile(r"-(?:\d{2,}-of-\d{2,}|0\d{3,})$") # Common on-disk split words -> the canonical concept they normalize to. _CANONICAL_ALIASES = { @@ -39,12 +42,6 @@ class ResolvedSplit: entries: list[FileEntry] -def _split_name(path: str) -> str: - """The shard-stripped file stem, e.g. train-00000-of-00003.parquet -> "train".""" - stem = PurePosixPath(path).name.split(".")[0] - return _SHARD_SUFFIX.sub("", stem) - - def _canonical_for(split_name: str) -> str | None: """Map a split name to its canonical concept, tolerating variant suffixes (train_prefs -> train).""" lowered = split_name.lower() @@ -54,12 +51,36 @@ def _canonical_for(split_name: str) -> str | None: return None +def is_split_directory(name: str) -> bool: + """Whether a directory name denotes a split rather than a partition. + + Partition grouping needs this to avoid turning ``train/`` and ``test/`` into two unrelated + partitions, each with its own schema and classification, when they are two splits of one dataset. + """ + return _canonical_for(name) is not None + + +def _split_name(path: str) -> str: + """The split a file belongs to. + + A split-named directory anywhere on the path wins over the filename, because the ``data/train/`` + layout names its shards ``0000.parquet`` — reading only the stem would file every split's shards + under the same meaningless name and collapse the whole dataset into one split. Nearest directory + to the file wins. Failing that, the shard-stripped stem: train-00000-of-00003.parquet -> "train". + """ + parts = PurePosixPath(path).parts + for directory in reversed(parts[:-1]): + if is_split_directory(directory): + return directory + return _SHARD_SUFFIX.sub("", parts[-1].split(".")[0]) + + def resolve_splits(entries: list[FileEntry]) -> list[ResolvedSplit]: """Group files into splits by path inference. - Each file's split name is its shard-stripped stem; the canonical concept is matched against - common aliases (val/valid/dev -> validation). When no file carries a recognizable split, every - file lands in one "default" split. + Each file's split name comes from a split-named directory on its path, else its shard-stripped + stem; the canonical concept is matched against common aliases (val/valid/dev -> validation). When + no file carries a recognizable split, every file lands in one "default" split. """ grouped: dict[str, list[FileEntry]] = {} for entry in entries: diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 6c9fc93208..7dfb00c424 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -3,6 +3,7 @@ """Tests for the profiling pipeline: digest, split/partition resolution, and envelope assembly.""" +import json from datetime import datetime, timezone import pyarrow as pa @@ -12,6 +13,7 @@ from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.pipeline import profile from nemo_datasets_plugin.profiler.splits import resolve_splits +from nemo_platform_plugin.files.dataset_profile import DatasetProfile FIXED_TIME = datetime(2026, 7, 13, 12, 0, 0, tzinfo=timezone.utc) @@ -62,6 +64,14 @@ def test_resolve_splits_normalizes_aliases(): assert splits == {"val": "validation", "dev": "validation"} +def test_resolve_splits_does_not_mistake_years_for_shard_numbers(): + # A bare trailing number only reads as a shard when it is zero-padded; otherwise dates and + # versions were being stripped, e.g. covid-19.jsonl -> a "covid" split. + assert [s.name for s in resolve_splits(_entries("covid-19.jsonl"))] == ["default"] + names = {s.name for s in resolve_splits(_entries("train-00000-of-00002.parquet", "data-2024.jsonl"))} + assert names == {"train", "data-2024"} + + def test_resolve_splits_falls_back_to_single_default(): splits = resolve_splits(_entries("shard-00000.parquet", "shard-00001.parquet")) assert len(splits) == 1 @@ -89,6 +99,26 @@ def test_group_partitions_splits_multiple_top_dirs(): assert [name for name, _ in parts] == ["main", "socratic"] +def test_group_partitions_does_not_treat_split_dirs_as_partitions(): + # train/ and test/ are one dataset's splits, not two datasets. + parts = group_partitions(_entries("train/data.parquet", "test/data.parquet")) + assert [name for name, _ in parts] == ["default"] + + +def test_resolve_splits_reads_the_split_directory(): + # The data// layout names every shard the same thing; only the directory carries + # the split, so reading the stem alone would collapse the dataset into one split. + splits = {s.name: s for s in resolve_splits(_entries("data/train/0000.parquet", "data/test/0000.parquet"))} + assert set(splits) == {"train", "test"} + assert splits["train"].canonical == "train" + assert splits["test"].canonical == "test" + + +def test_resolve_splits_prefers_directory_over_stem(): + splits = resolve_splits(_entries("main/train-00000-of-00001.parquet")) + assert [s.name for s in splits] == ["train"] # no split dir on the path, so the stem is used + + # --- end-to-end profile() ------------------------------------------------------------------------ @@ -119,8 +149,11 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert partition.stats["prompt"].text is not None assert partition.classification.dataset_type == "prompt_only" # a lone prompt column, no target + # strategy is the policy, exhaustive is the outcome: a capped run over files that all fit under + # the cap is still a full scan, which is why the contract keeps the two fields independent. + assert result.sampling.strategy == "head_per_file" assert result.sampling.exhaustive is True - assert result.sampling.strategy == "full" + assert result.sampling.per_file_row_cap == 1000 assert result.sampling.rows_scanned == 3 assert result.sampling.rows_total == 3 assert result.sampling.files_scanned == 2 @@ -147,6 +180,36 @@ def test_profile_multiple_directories_become_partitions(tmp_path): assert all(p.file_format == "parquet" for p in result.partitions) +def test_profile_top_level_split_dirs_become_one_partition(tmp_path): + # train/ + test/ is one dataset with two splits, not two datasets. As separate partitions each + # would derive its own schema and classification, and the split structure would disappear. + _write_parquet(tmp_path / "train" / "data.parquet", [{"prompt": "a", "completion": "b"}]) + _write_parquet(tmp_path / "test" / "data.parquet", [{"prompt": "c", "completion": "d"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert [p.name for p in result.partitions] == ["default"] + splits = {s.name: s for s in result.partitions[0].splits} + assert set(splits) == {"train", "test"} + assert splits["train"].canonical == "train" + assert splits["test"].num_examples == 1 + + +def test_profile_nested_split_dirs_keep_splits_apart(tmp_path): + # data// shards are all named alike, so only the directory distinguishes them. + # Reading the stem alone pooled train and test into a single "default" split. + _write_parquet(tmp_path / "data" / "train" / "0000.parquet", [{"prompt": "a"}, {"prompt": "b"}]) + _write_parquet(tmp_path / "data" / "test" / "0000.parquet", [{"prompt": "c"}]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert [p.name for p in result.partitions] == ["default"] + splits = {s.name: s for s in result.partitions[0].splits} + assert set(splits) == {"train", "test"} + assert splits["train"].num_examples == 2 + assert splits["test"].num_examples == 1 + + def test_profile_splits_mixed_formats_into_separate_partitions(tmp_path): # A directory holding two formats must not profile as one partition: features/stats would be # derived from one format's schema but measured over rows from both. Each format becomes its own @@ -167,6 +230,52 @@ def test_profile_splits_mixed_formats_into_separate_partitions(tmp_path): assert result.sampling.exhaustive is True +def test_profile_unions_columns_across_shards(tmp_path): + # A column that appears only in a later shard must still reach features/stats. Taking the first + # shard's schema would drop it entirely. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a", "completion": "b"}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"prompt": "c", "completion": "d", "score": 3}]) + + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert [f.name for f in partition.features] == ["prompt", "completion", "score"] + assert partition.stats["score"].null_rate == 0.5 # absent in the first shard, and said so + + +def test_profile_is_invariant_to_shard_order(tmp_path, tmp_path_factory): + # The same rows must profile the same way regardless of which shard sorts first. First-wins + # schema selection made this data classify as prompt_completion or scored_response depending + # purely on filename order. + narrow = [{"prompt": "a", "completion": "b"}] + wide = [{"prompt": "c", "completion": "d", "score": 3}] + + forward = tmp_path_factory.mktemp("forward") + _write_parquet(forward / "train-00000-of-00002.parquet", narrow) + _write_parquet(forward / "train-00001-of-00002.parquet", wide) + + reverse = tmp_path_factory.mktemp("reverse") + _write_parquet(reverse / "train-00000-of-00002.parquet", wide) + _write_parquet(reverse / "train-00001-of-00002.parquet", narrow) + + first = profile(LocalFileSource(forward), created_at=FIXED_TIME).partitions[0] + second = profile(LocalFileSource(reverse), created_at=FIXED_TIME).partitions[0] + + assert [(f.name, f.dtype) for f in first.features] == [(f.name, f.dtype) for f in second.features] + assert first.classification.dataset_type == second.classification.dataset_type == "scored_response" + + +def test_profile_survives_conflicting_shard_schemas(tmp_path): + # Two shards disagreeing on a column's type has no right answer at the schema level; fall back to + # inferring from the rows (which widens to json) rather than asserting one shard over the other. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"score": 1}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"score": "high"}]) + + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert [f.name for f in partition.features] == ["score"] + assert partition.features[0].dtype == "json" # mixed, and honest about it + + def test_profile_isolates_unreadable_files(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) (tmp_path / "test-00000-of-00001.parquet").write_bytes(b"not a real parquet file") @@ -177,8 +286,84 @@ def test_profile_isolates_unreadable_files(tmp_path): assert splits["train"].num_examples == 1 assert splits["test"].num_examples is None # unreadable -> count unknown, not a crash assert splits["test"].files[0].num_rows is None + assert splits["test"].files[0].error is not None # ...and the profile says why assert result.sampling.exhaustive is False # a file could not be fully parsed assert result.sampling.rows_total is None + assert result.sampling.files_scanned == 1 # one file was actually read; the other never opened + + +def test_profile_row_cap_bounds_reads_and_says_so(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": i} for i in range(10)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=4) + + assert result.sampling.rows_scanned == 4 + assert result.sampling.per_file_row_cap == 4 + assert result.sampling.exhaustive is False # 4 of 10 rows is not a full scan + assert result.sampling.rows_total is None + assert result.partitions[0].splits[0].num_examples == 10 # the footer count survives sampling + + +def test_profile_uncapped_read_is_a_full_scan(tmp_path): + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": i} for i in range(10)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=None) + + assert result.sampling.strategy == "full" + assert result.sampling.per_file_row_cap is None + assert result.sampling.exhaustive is True + assert result.sampling.rows_scanned == 10 + + +def test_profile_cap_larger_than_a_jsonl_file_keeps_it_exhaustive(tmp_path): + # jsonl has no footer, so a cap could easily cost the exact count on files that never hit it. + # Reading to EOF under the cap must stay exact, or capping would degrade every small dataset. + (tmp_path / "train.jsonl").write_text('{"a": 1}\n{"a": 2}\n') + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=1000) + + assert result.partitions[0].splits[0].num_examples == 2 + assert result.sampling.exhaustive is True + assert result.sampling.rows_total == 2 + + +def test_profile_reports_unsupported_data_files(tmp_path): + # A directory of formats we cannot read must not profile as an exhaustively scanned empty + # dataset — that is indistinguishable from a dataset that really is empty. + (tmp_path / "train.csv").write_text("prompt,completion\na,b\n") + (tmp_path / "test.arrow").write_bytes(b"\x00") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.partitions == [] + assert result.sampling.exhaustive is False # we scanned nothing, and admit it + assert result.sampling.rows_total is None # not 0: "empty" would be a lie + assert result.profiler_info["unsupported_files"] == ["test.arrow", "train.csv"] + + +def test_profile_ignores_non_data_files_without_penalty(tmp_path): + # A README is genuinely not data, so it must not cost exhaustiveness the way a .csv does. + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + (tmp_path / "README.md").write_text("a dataset card") + (tmp_path / "LICENSE").write_text("apache") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.sampling.exhaustive is True + assert "unsupported_files" not in result.profiler_info + + +def test_profile_records_a_partial_jsonl_read(tmp_path): + # One corrupt line costs that line, not the file — but the profile must still say the file was + # only partly understood, rather than presenting a clean-looking count. + (tmp_path / "train.jsonl").write_text('{"a": 1}\n{"a": 2\n{"a": 3}\n') + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + record = result.partitions[0].splits[0].files[0] + assert record.num_rows == 2 # the readable rows survived + assert record.error is not None and "line 2" in record.error + assert result.sampling.exhaustive is False # a line was lost, so this is not a full scan def test_profile_classifies_roles_type_and_verifiability(tmp_path): @@ -194,6 +379,38 @@ def test_profile_classifies_roles_type_and_verifiability(tmp_path): assert partition.classification.verifiability.coverage == 1.0 +def test_profile_sharegpt_dataset_is_a_chat_dataset(tmp_path): + # End to end: {from, value} must reach the messages dtype, carry stats, and classify as chat + # rather than falling through to `unknown` with nothing measured. + conversation = [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello"}] + (tmp_path / "train.jsonl").write_text(json.dumps({"conversations": conversation}) + "\n") + + partition = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert partition.features[0].dtype == "messages" + assert partition.features[0].semantic_role == "messages" + assert partition.classification.dataset_type == "messages" + assert partition.stats["conversations"].messages.roles_seen == ["human", "gpt"] + + +def test_profile_degrades_one_partition_when_measurement_fails(tmp_path, monkeypatch): + # Reads are isolated per file, but schema/stats/classification ran unguarded, so one odd value + # could abort an otherwise complete profile. Structure must survive a measurement failure. + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) + monkeypatch.setattr(pipeline_module, "derive_stats", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise + + partition = result.partitions[0] + assert partition.splits[0].num_examples == 2 # structure survives + assert partition.stats == {} + assert partition.classification.dataset_type == "unknown" + assert [e.kind for e in partition.classification.evidence] == ["error"] + assert "RuntimeError" in partition.classification.evidence[0].detail # says what failed + + def test_profile_is_deterministic(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) source = LocalFileSource(tmp_path) @@ -212,6 +429,45 @@ def test_profile_tolerates_non_object_jsonl_lines(tmp_path): assert result.sampling.exhaustive is True +def test_profile_survives_a_hostile_directory(tmp_path): + """Everything that can go wrong at once must still yield a profile that says what went wrong. + + Each of these individually used to either abort the run or vanish silently; this is the shape of + bug that got through, so it is worth asserting as one scenario rather than only in isolation. + """ + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a", "completion": "b"}]) + (tmp_path / "train-00001-of-00002.parquet").write_bytes(b"not a parquet file") # corrupt + (tmp_path / "extra.jsonl").write_text( + '{"messages": [{"role": 1, "content": "hi"}]}\n' # non-string role + '{"messages": [{"role": "user"\n' # truncated line + "[1, 2, 3]\n" # valid JSON, not a row + '{"messages": [{"role": "user", "content": "ok"}]}\n' + ) + (tmp_path / "leftovers.csv").write_text("a,b\n1,2\n") # recognizable data, no reader + (tmp_path / "README.md").write_text("a dataset card") # not data at all + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise + + # Nothing here is exhaustive, and the profile says so rather than looking clean. + assert result.sampling.exhaustive is False + assert result.sampling.rows_total is None + assert result.profiler_info["unsupported_files"] == ["leftovers.csv"] + + records = {f.path: f for p in result.partitions for s in p.splits for f in s.files} + assert records["train-00001-of-00002.parquet"].error is not None # corrupt file, named and explained + assert records["extra.jsonl"].error is not None # partial parse, named and explained + + # The readable parquet rows still produced a real classification. + parquet_partition = next(p for p in result.partitions if p.file_format == "parquet") + assert parquet_partition.classification.dataset_type == "prompt_completion" + # ...and the odd jsonl rows were measured rather than aborting the run. + jsonl_partition = next(p for p in result.partitions if p.file_format == "jsonl") + assert jsonl_partition.stats["messages"].messages.roles_seen == ["1", "user"] + + # The whole thing still round-trips as a stored profile. + assert DatasetProfile.model_validate_json(result.model_dump_json()) == result + + def test_profile_digest_covers_only_stored_files(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) without_card = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) From 118b103621722292b2689a5e6e9765d1b0141649 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Tue, 4 Aug 2026 16:47:42 -0400 Subject: [PATCH 19/44] fix(datasets): measure content probes independently of role assignment Content detection was gated behind name detection. _detect_verifiability reached the rows through _completion_texts, which needed a `completion` role to know which column to read, so a dataset whose columns are named `q`/`a` rather than `question`/`answer` lost its entire classification layer -- roles, type, format, prompt_form and verifiability -- even though the `#### ` markers were sitting in the data and the regex would have matched them. Move the probes into stats as derive_probes(), measured over every column, and leave classify to interpret the counts. Roles still order that interpretation: a column known to be the ground truth beats one that merely looks like it, and a named completion is still the authoritative place to look. They just no longer gate it, and the finding names the column it came from. That also fixes ShareGPT-shaped data losing verifiability outright. A chat column's probe text is read through _message_field, so `{from, value}` behaves like `{role, content}`, as it already does in schema derivation and message stats. Second fix in the same stage: stop trusting a partial declared schema. derive_features uses the arrow schema if present *at all* and ignores the rows, so a group where only some files declare one erased every column the schemaless files were the sole witness for. _measure now infers from rows unless every file that contributed rows declared a schema. This is the defect _split_by_format worked around by forcing partitions to be format-homogeneous; fixing it where it lives is what lets that grouping go. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/classify.py | 139 ++++++++++-------- .../nemo_datasets_plugin/profiler/pipeline.py | 35 ++++- .../nemo_datasets_plugin/profiler/stats.py | 75 +++++++++- plugins/nemo-datasets/tests/test_classify.py | 43 ++++++ plugins/nemo-datasets/tests/test_pipeline.py | 18 ++- plugins/nemo-datasets/tests/test_stats.py | 50 ++++++- 6 files changed, 286 insertions(+), 74 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index 915cf33e92..207a8400a3 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -5,13 +5,17 @@ Roles are inferred from column names, gated by dtype, and stacked onto the feature nodes as ``semantic_role`` markers. The dataset type is the most specific structure the assigned roles -satisfy. Verifiability and content-probe corroboration are added by a later stage. +satisfy. + +Content probes are *measured* in :mod:`stats` over every column; this module only interprets the +counts. Roles still order that interpretation — a column known to be the ground truth is a better +answer than one that merely looks like it — but they no longer gate it, so a dataset whose columns +carry unrecognized names keeps whatever its content proves. """ from __future__ import annotations -import re - +from nemo_datasets_plugin.profiler.stats import ColumnProbes, derive_probes from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, Evidence, @@ -207,11 +211,7 @@ def has(*required: str) -> bool: return "unknown" -# --- content probes ------------------------------------------------------------------------------ - -_TRANSCRIPT_MARKER = re.compile(r"\n\n(?:Human|Assistant|User):") -_GSM8K_ANSWER = re.compile(r"####\s*-?[\d.,/]+\s*$") -_BOXED_ANSWER = re.compile(r"\\boxed\{") +# --- interpreting the content probes -------------------------------------------------------------- # A verification target must cover at least this fraction of sampled rows to be asserted. Below it, # a "hit" is noise -- e.g. one completion in thousands coincidentally ending in `#### ` does @@ -223,51 +223,51 @@ def _pct(fraction: float) -> str: return f"{round(fraction * 100)}%" -def _completion_texts(features: list[FeatureSchema], rows: list[dict]) -> list[str]: - completion = next((feature for feature in features if feature.semantic_role == "completion"), None) - if completion is None: - return [] - texts: list[str] = [] - for row in rows: - value = row.get(completion.name) - if isinstance(value, str): - texts.append(value) - elif isinstance(value, list) and value and isinstance(value[-1], dict): - content = value[-1].get("content") # the final assistant turn for a conversational completion - if isinstance(content, str): - texts.append(content) - return texts - - -def _detect_verifiability(features: list[FeatureSchema], rows: list[dict]) -> Verifiability | None: - if not rows: - return None +def _detect_verifiability(features: list[FeatureSchema], probes: dict[str, ColumnProbes]) -> Verifiability | None: + """The strongest verification target the probes found, if any clears the coverage floor. - # Each method wins only if it clears the coverage floor; otherwise fall through to the next, so a - # sparse ground_truth column can still yield to an extractable-answer signal instead of masking it. + Each method wins only if it clears the floor; otherwise fall through to the next, so a sparse + ground_truth column yields to an extractable-answer signal instead of masking it. + """ ground_truth = next((feature for feature in features if feature.semantic_role == "ground_truth"), None) if ground_truth is not None: - present = sum(1 for row in rows if row.get(ground_truth.name) not in (None, "", [])) - coverage = present / len(rows) - if coverage >= _MIN_VERIFIABILITY_COVERAGE: - detail = f"'{ground_truth.name}' present in {_pct(coverage)} of {len(rows)} sampled rows" - return Verifiability( - method="ground_truth_column", - coverage=coverage, - evidence=[Evidence(kind="content_probe", detail=detail)], - ) - - texts = _completion_texts(features, rows) - if texts: - hits = sum(1 for text in texts if _GSM8K_ANSWER.search(text) or _BOXED_ANSWER.search(text)) - coverage = hits / len(texts) - if coverage >= _MIN_VERIFIABILITY_COVERAGE: - detail = f"completion ends with an extractable answer (#### or \\boxed) in {_pct(coverage)} of {len(texts)} sampled rows" - return Verifiability( - method="extractable_final_answer", - coverage=coverage, - evidence=[Evidence(kind="content_probe", detail=detail)], - ) + probe = probes.get(ground_truth.name) + if probe is not None and probe.rows: + coverage = probe.non_empty / probe.rows + if coverage >= _MIN_VERIFIABILITY_COVERAGE: + detail = f"'{ground_truth.name}' present in {_pct(coverage)} of {probe.rows} sampled rows" + return Verifiability( + method="ground_truth_column", + coverage=coverage, + evidence=[Evidence(kind="content_probe", detail=detail)], + ) + + # A named completion is the authoritative place to look. Without one, take whichever column the + # probes found the strongest signal in and name it — the markers are a fact about that column + # whether or not its name happened to be in the alias table. + completion = next((feature for feature in features if feature.semantic_role == "completion"), None) + searched = [completion] if completion is not None else features + best_name: str | None = None + best_coverage = 0.0 + for feature in searched: + probe = probes.get(feature.name) + if probe is None or not probe.texts: + continue + coverage = probe.extractable_answer / probe.texts + if coverage > best_coverage: + best_name, best_coverage = feature.name, coverage + + if best_name is not None and best_coverage >= _MIN_VERIFIABILITY_COVERAGE: + sampled = probes[best_name].texts + detail = ( + f"'{best_name}' ends with an extractable answer (#### or \\boxed) in " + f"{_pct(best_coverage)} of {sampled} sampled rows" + ) + return Verifiability( + method="extractable_final_answer", + coverage=best_coverage, + evidence=[Evidence(kind="content_probe", detail=detail)], + ) return None @@ -279,18 +279,22 @@ def _common_prefix_len(left: str, right: str) -> int: return index -def _implicit_prompt_evidence(features: list[FeatureSchema], rows: list[dict]) -> Evidence | None: +def _implicit_prompt_evidence( + features: list[FeatureSchema], probes: dict[str, ColumnProbes], rows: list[dict] +) -> Evidence | None: + targets = [f for f in features if f.semantic_role in {"chosen", "rejected", "completion"} and f.dtype == "string"] + counted = [probes[f.name] for f in targets if f.name in probes] + sampled = sum(probe.texts for probe in counted) + marked = sum(probe.transcript_marker for probe in counted) + if sampled and marked: + detail = f"embedded transcript markers in {_pct(marked / sampled)} of sampled completions - prompt is embedded" + return Evidence(kind="content_probe", detail=detail) + + # The shared-prefix check is *relational* — it compares two columns against each other — so it + # has no per-column probe to read and still works from the rows themselves. if not rows: return None - targets = [f for f in features if f.semantic_role in {"chosen", "rejected", "completion"} and f.dtype == "string"] - texts = [value for f in targets for row in rows if isinstance((value := row.get(f.name)), str)] - if texts: - marked = sum(1 for text in texts if _TRANSCRIPT_MARKER.search(text)) - if marked: - detail = f"embedded transcript markers in {_pct(marked / len(texts))} of sampled completions - prompt is embedded" - return Evidence(kind="content_probe", detail=detail) - chosen = next((f for f in features if f.semantic_role == "chosen" and f.dtype == "string"), None) rejected = next((f for f in features if f.semantic_role == "rejected" and f.dtype == "string"), None) if chosen is not None and rejected is not None: @@ -309,14 +313,21 @@ def _implicit_prompt_evidence(features: list[FeatureSchema], rows: list[dict]) - def classify( - features: list[FeatureSchema], stats: dict[str, ColumnStats], rows: list[dict] | None = None + features: list[FeatureSchema], + stats: dict[str, ColumnStats], + rows: list[dict] | None = None, + *, + probes: dict[str, ColumnProbes] | None = None, ) -> PartitionClassification: """Assign roles onto ``features`` in place and return the partition's classification. - ``rows`` (the sampled rows) drive the content probes — verifiability and implicit-prompt - detection; role/axis/type inference needs only the schema and stats. + ``probes`` are the per-column content measurements from :func:`~.stats.derive_probes`. They are + a pure function of ``(features, rows)``, so a caller that has not already computed them can pass + ``rows`` alone and get them derived here; the pipeline passes them in to avoid the second pass. + Role/axis/type inference needs neither — only the schema and stats. """ rows = rows or [] + probes = derive_probes(features, rows) if probes is None else probes _assign_roles(features, stats) roles = {feature.semantic_role for feature in features if feature.semantic_role} dataset_type = _detect_type(features, stats) @@ -330,7 +341,7 @@ def classify( if fmt is not None: evidence.append(Evidence(kind="column_dtype", detail=f"{fmt} format from role column dtypes")) if prompt_form == "implicit": - embedded = _implicit_prompt_evidence(features, rows) + embedded = _implicit_prompt_evidence(features, probes, rows) if embedded is not None: evidence.append(embedded) @@ -339,6 +350,6 @@ def classify( dataset_type=dataset_type, format=fmt, prompt_form=prompt_form, - verifiability=_detect_verifiability(features, rows), + verifiability=_detect_verifiability(features, probes), evidence=evidence, ) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 61c2a0d1a0..2eef5bf50b 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -28,7 +28,7 @@ from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import resolve_splits -from nemo_datasets_plugin.profiler.stats import derive_stats +from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, DatasetProfile, @@ -172,21 +172,35 @@ def _unify_schemas(schemas: list[pa.Schema]) -> pa.Schema | None: def _measure( - partition_rows: list[dict], arrow_schemas: list[pa.Schema], exhaustive: bool + partition_rows: list[dict], + arrow_schemas: list[pa.Schema], + *, + exhaustive: bool, + all_declared: bool, ) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: """Derive schema, stats and classification, degrading to structure-only if any of it fails. - These three stages are pure computation over rows already in memory, so a failure here is either - a profiler bug or data shaped in a way no detector anticipated. Reads are already isolated per + ``all_declared`` says whether *every* file that contributed rows carried a declared schema. When + one did not, the unified schema describes only some of the rows, and using it would erase any + column the schemaless files were the sole witness for — so infer from the rows instead, which + sees all of them. Declared type fidelity (int32 widening to int64) is the cost, and it is the + honest one: a declared schema cannot be asserted over files that declare nothing. + + These stages are pure computation over rows already in memory, so a failure here is either a + profiler bug or data shaped in a way no detector anticipated. Reads are already isolated per file; leaving this stage unguarded meant one odd value — a chat message whose ``role`` is a number — could abort an otherwise complete profile from the one place nothing was catching. The partition's structure (files, splits, row counts) is established by then and stays useful, so the failure costs its measurements and says so, rather than the entire run. """ try: - features = derive_features(partition_rows, _unify_schemas(arrow_schemas)) + declared = _unify_schemas(arrow_schemas) if all_declared else None + features = derive_features(partition_rows, declared) stats = derive_stats(features, partition_rows, exhaustive=exhaustive) - return features, stats, classify(features, stats, partition_rows) + # Probes are measured over every column, independent of the roles classify is about to + # assign, so a content signal survives a column name the alias table does not know. + probes = derive_probes(features, partition_rows) + return features, stats, classify(features, stats, partition_rows, probes=probes) except Exception as exc: detail = f"could not measure this partition: {type(exc).__name__}: {exc}" return [], {}, PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]) @@ -213,6 +227,7 @@ def _profile_partition( """ partition_rows: list[dict] = [] arrow_schemas: list[pa.Schema] = [] + all_declared = True # every file that contributed rows carried a declared schema rows_scanned = 0 files_read = 0 partition_scanned = True @@ -243,6 +258,10 @@ def _profile_partition( partition_rows.extend(result.rows) if result.arrow_schema is not None: arrow_schemas.append(result.arrow_schema) + elif result.rows: + # Rows with no schema behind them: the unified schema no longer covers the + # partition, so _measure must infer from rows rather than trust a partial one. + all_declared = False # Exhaustive requires parsing every row; a known footer count alone is not enough, and # a partial read (corrupt lines skipped) is not exhaustive however many rows it got. scanned_all = num_rows is not None and result.rows_scanned >= num_rows and error is None @@ -270,7 +289,9 @@ def _profile_partition( num_examples=split_examples if split_counts_known else None, ) ) - features, stats, classification = _measure(partition_rows, arrow_schemas, partition_scanned) + features, stats, classification = _measure( + partition_rows, arrow_schemas, exhaustive=partition_scanned, all_declared=all_declared + ) partition = PartitionProfile( name=name, file_format=file_format, diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 1b97eea3e6..2d3a7b115b 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -1,19 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Per-column statistics. +"""Per-column statistics and content probes. Given a partition's features and its sampled rows, measure each top-level column according to its dtype: length quantiles and corruption signals for text, min/max/mean for numbers, chat-shape signals for messages, and cardinality for both. The result is sparse — a column with nothing worth measuring is omitted. Row values themselves are never stored, except a proven small enumeration under ``categorical.values`` when the read was exhaustive. + +:func:`derive_probes` additionally reads each column's *content* — answer markers, embedded +transcripts — as plain per-column counts. Those are measurements, not interpretations: what they +mean is classification's job, and keeping the looking here is what stops a content signal from +being reachable only through a correctly named column. """ from __future__ import annotations import math import re +from dataclasses import dataclass from typing import Any from nemo_platform_plugin.files.dataset_profile import ( @@ -256,3 +262,70 @@ def _valid_alternation(messages: list) -> bool: """True when user/assistant turns alternate (ignoring any leading system turns).""" roles = [_role_of(m) for m in messages if isinstance(m, dict) and _role_of(m) != "system"] return all(roles[i] != roles[i + 1] for i in range(len(roles) - 1)) + + +# --- content probes ------------------------------------------------------------------------------ + +# Probes run over *every* column, not only role-assigned ones. Gating them on roles made a content +# signal reachable only through a recognized column name: a dataset whose answer column is called +# `a` instead of `answer` lost verifiability entirely, even though the markers were sitting in the +# data and the regex would have matched them. Classification reads these counts and decides what +# they mean; it no longer does the looking. +_TRANSCRIPT_MARKER = re.compile(r"\n\n(?:Human|Assistant|User):") +_GSM8K_ANSWER = re.compile(r"####\s*-?[\d.,/]+\s*$") +_BOXED_ANSWER = re.compile(r"\\boxed\{") + + +@dataclass(frozen=True) +class ColumnProbes: + """What the content probes saw in one column across the sampled rows. + + Internal to the profiler rather than part of the stored contract: these are inputs to + classification, and promoting them to durable per-column facts is a separate contract change. + Counts, not rates — the caller divides, so a zero denominator stays visible instead of becoming + a silent 0.0. + """ + + rows: int # rows considered for this column + non_empty: int # value present and not "" / [] / {} — a usable target of any dtype + texts: int # rows that yielded text: a string, or a chat column's final turn + extractable_answer: int # of `texts`, how many carry `#### ` or `\boxed{` + transcript_marker: int # of `texts`, how many embed a Human:/Assistant: transcript + + +def derive_probes(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> dict[str, ColumnProbes]: + """Run the content probes over every top-level column, keyed by column name.""" + probes: dict[str, ColumnProbes] = {} + for feature in features: + # Duplicate parquet field names: first wins, matching derive_stats so the two agree on which. + if feature.name in probes: + continue + probes[feature.name] = _column_probes([row.get(feature.name) for row in rows]) + return probes + + +def _column_probes(values: list[Any]) -> ColumnProbes: + texts = [text for value in values if (text := _probe_text(value)) is not None] + return ColumnProbes( + rows=len(values), + non_empty=sum(1 for value in values if value not in (None, "", [], {})), + texts=len(texts), + extractable_answer=sum(1 for text in texts if _GSM8K_ANSWER.search(text) or _BOXED_ANSWER.search(text)), + transcript_marker=sum(1 for text in texts if _TRANSCRIPT_MARKER.search(text)), + ) + + +def _probe_text(value: Any) -> str | None: + """The text a probe reads from one cell: the string itself, or a chat column's final turn. + + The final turn is read through :func:`_message_field`, so ShareGPT's ``{from, value}`` spelling + works like ``{role, content}``. Both are handled everywhere else in this module and in schema + derivation; missing it here cost every ShareGPT-shaped dataset its verifiability. + """ + if isinstance(value, str): + return value + if isinstance(value, list) and value and isinstance(value[-1], dict): + content = _message_field(value[-1], "content", "value") + if isinstance(content, str): + return content + return None diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index 6035e632c8..9c84a2174e 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -4,6 +4,7 @@ """Tests for classification: role assignment, format/prompt-form axes, and dataset type.""" from nemo_datasets_plugin.profiler.classify import classify +from nemo_datasets_plugin.profiler.stats import derive_probes from nemo_platform_plugin.files.dataset_profile import ( CategoricalStats, ColumnStats, @@ -252,3 +253,45 @@ def test_bare_scalar_ground_truth_alias_is_still_rejected(): features = [_f("ground_truth", "int64")] classify(features, {}) assert features[0].semantic_role is None + + +def test_verifiability_survives_an_unrecognized_column_name(): + # Gating the probes on roles made a content signal reachable only through a recognized column + # name: the `#### ` markers were in the data and the regex would have matched them, but + # nothing knew where to look. The finding must name the column it came from. + features = [_f("q", "string"), _f("a", "string")] + rows = [{"q": "what is 2+2?", "a": f"add them #### {i}"} for i in range(10)] + result = classify(features, {}, rows) + + assert {f.semantic_role for f in features} == {None} # still unroled, and honest about it + assert result.dataset_type == "unknown" + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + assert "'a'" in result.verifiability.evidence[0].detail + + +def test_a_named_completion_still_decides_where_to_look(): + # Roles order the interpretation even though they no longer gate it: a column *known* to be the + # completion is a better answer than one that merely looks like it. + features = [_f("completion", "string"), _f("notes", "string")] + rows = [{"completion": "just prose", "notes": "scratch #### 9"} for _ in range(10)] + assert classify(features, {}, rows).verifiability is None + + +def test_verifiability_reads_a_sharegpt_conversational_completion(): + features = [_f("prompt", "string"), _f("completion", "messages")] + rows = [{"prompt": "q", "completion": [{"from": "human", "value": "q"}, {"from": "gpt", "value": "#### 4"}]}] + result = classify(features, {}, rows) + assert result.verifiability.method == "extractable_final_answer" + assert result.verifiability.coverage == 1.0 + + +def test_precomputed_probes_and_derived_probes_agree(): + # classify() derives probes when it is not given them; the pipeline passes them in to avoid the + # second pass. The two paths must not drift. + rows = [{"prompt": "q", "completion": f"steps #### {i}"} for i in range(10)] + derived = classify([_f("prompt", "string"), _f("completion", "string")], {}, rows) + features = [_f("prompt", "string"), _f("completion", "string")] + passed_in = classify(features, {}, rows, probes=derive_probes(features, rows)) + assert derived.verifiability.coverage == passed_in.verifiability.coverage + assert derived.verifiability.evidence[0].detail == passed_in.verifiability.evidence[0].detail diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 7dfb00c424..e5f52f0a44 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -11,7 +11,7 @@ from nemo_datasets_plugin.profiler.digest import content_digest from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource from nemo_datasets_plugin.profiler.partition import group_partitions -from nemo_datasets_plugin.profiler.pipeline import profile +from nemo_datasets_plugin.profiler.pipeline import _measure, profile from nemo_datasets_plugin.profiler.splits import resolve_splits from nemo_platform_plugin.files.dataset_profile import DatasetProfile @@ -501,3 +501,19 @@ def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): records = {f.path: f for p in result.partitions for s in p.splits for f in s.files} assert records["extra.xyz"].num_rows is None # kept, but unreadable assert result.sampling.exhaustive is False + + +def test_measure_infers_from_rows_when_some_files_declared_no_schema(): + # derive_features uses the declared schema *if present at all*, so a group where only some files + # declare one erased every column the schemaless files were the sole witness for. That is the + # defect `_split_by_format` worked around by making partitions format-homogeneous; the fix + # belongs in schema derivation, and dropping that homogeneity is what makes this path reachable. + declared = pa.schema([pa.field("prompt", pa.string())]) + rows = [{"prompt": "a"}, {"prompt": "b", "extra": "only in the schemaless file"}] + + features, stats, _ = _measure(rows, [declared], exhaustive=True, all_declared=False) + assert [f.name for f in features] == ["prompt", "extra"] # the sole witness survives + assert set(stats) <= {f.name for f in features} + + features, _, _ = _measure(rows, [declared], exhaustive=True, all_declared=True) + assert [f.name for f in features] == ["prompt"] # declared schema trusted when it covers everything diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 4599cce5db..b764909f61 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -3,7 +3,7 @@ """Tests for per-column statistics.""" -from nemo_datasets_plugin.profiler.stats import derive_stats +from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema @@ -161,3 +161,51 @@ def test_unmeasured_dtypes_are_omitted(): def test_null_rate_is_reported(): stats = derive_stats([_feature("t", "string")], _rows("t", ["a", None, "c", None]), exhaustive=False)["t"] assert stats.null_rate == 0.5 + + +# --- content probes ------------------------------------------------------------------------------ + + +def test_probes_are_measured_for_every_column_not_just_named_ones(): + # The whole point of measuring probes here rather than in classify: a column whose name the + # alias table does not know still gets its content read. + features = [_feature("q", "string"), _feature("a", "string")] + rows = [{"q": "what is 2+2?", "a": "add them #### 4"}, {"q": "and 3+3?", "a": "no final answer"}] + probes = derive_probes(features, rows) + + assert set(probes) == {"q", "a"} + assert probes["a"].texts == 2 + assert probes["a"].extractable_answer == 1 + assert probes["q"].extractable_answer == 0 + + +def test_probes_read_the_final_turn_of_a_chat_column(): + rows = [{"m": [{"role": "user", "content": "q"}, {"role": "assistant", "content": "steps #### 7"}]}] + probes = derive_probes([_feature("m", "messages")], rows) + assert probes["m"].texts == 1 + assert probes["m"].extractable_answer == 1 + + +def test_probes_read_the_sharegpt_message_spelling(): + # {from, value} is handled in schema derivation and message stats; reading only {role, content} + # here cost every ShareGPT-shaped dataset its verifiability. + rows = [{"m": [{"from": "human", "value": "q"}, {"from": "gpt", "value": "steps #### 7"}]}] + probes = derive_probes([_feature("m", "messages")], rows) + assert probes["m"].texts == 1 + assert probes["m"].extractable_answer == 1 + + +def test_probes_count_non_empty_across_container_dtypes(): + # `non_empty` is what a ground_truth column's coverage is measured from, and a verification + # target is just as often a list or struct as a string. + features = [_feature("gt", "list")] + rows = [{"gt": [{"in": "1"}]}, {"gt": []}, {"gt": None}] + probes = derive_probes(features, rows) + assert probes["gt"].rows == 3 + assert probes["gt"].non_empty == 1 + + +def test_probes_detect_embedded_transcripts(): + rows = [{"c": "\n\nHuman: hi\n\nAssistant: hello"}, {"c": "plain prose"}] + probes = derive_probes([_feature("c", "string")], rows) + assert probes["c"].transcript_marker == 1 From 0893185338db49fbe73d9b9e897555be16ed7af2 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Tue, 4 Aug 2026 16:54:11 -0400 Subject: [PATCH 20/44] refactor(files): drop the profile content digest rather than repair it The digest could never match the staleness check it existed to feed. The profiler hashed only `data_entries` (recognized extensions); the Files service hashed every file the listing returned. A fileset holding train.parquet and a README -- which is to say almost every real one -- mismatched on the first comparison and reported `stale` forever, immediately after a successful profile. Neither side's tests caught it: each asserted its own rule, and the Files tests built the expected value by calling the digest on their own data-file-only fixture. Repairing it would mean picking one file set, and that is the deeper problem: a stored digest freezes "which files count as inputs" into the data at write time. That judgment moves -- once card front-matter drives split declaration, README.md becomes an input -- and changing it would invalidate every stored profile at once, with no way to tell a real change from a definition change. The digest was never load-bearing anyway. The FileRecords already are the input manifest, and split membership is exhaustive and disjoint, so staleness needs a fresh listing plus a comparison against data the profile already holds. The listing is the entire cost; comparing 10k records against 10k listings is microseconds either way, and a direct comparison answers "3 shards added, 1 removed" instead of "different". (path, size) could not see a same-size in-place edit regardless. So: remove it. Profiling is user-triggered today, so nothing is waiting on staleness, and adding a field back later is a minor bump that stored profiles tolerate -- which the new test pins, since profiles already written with the field must keep loading. The digest test that asserted the broken rule is replaced by one asserting the invariant underneath it: the stored FileRecords reproduce the partition's input list exactly, which is what makes a listing comparison possible at all. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 23 +++++--- .../tests/files/test_dataset_profile.py | 16 ++++-- .../nemo_datasets_plugin/profiler/digest.py | 28 ---------- .../nemo_datasets_plugin/profiler/pipeline.py | 3 -- plugins/nemo-datasets/tests/test_pipeline.py | 54 +++++++------------ 5 files changed, 47 insertions(+), 77 deletions(-) delete mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 3ea4e69f20..43e676a5bd 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -249,8 +249,10 @@ class ColumnStats(BaseModel): class FileRecord(BaseModel): """One physical file, measured. - Stores the exact digest inputs (so a profile self-describes its ``content_digest`` and per-file - staleness is computable) plus what the reader learned cheaply. + Stores the file's identity as the listing reported it — path, size, checksum — plus what the + reader learned cheaply. Concatenating ``files`` across a partition's splits reproduces that + partition's input list exactly, which is what lets a consumer compare a stored profile against a + fresh listing to see what changed. """ path: str = Field(description="Relative path within the fileset.") @@ -258,8 +260,9 @@ class FileRecord(BaseModel): checksum: str | None = Field( default=None, description=( - 'As the Files service reports it (e.g. "sha256:..."). None falls back to a (path, size) digest, ' - "which cannot detect a same-size in-place edit." + 'As the Files service reports it (e.g. "sha256:..."), when it reports one at all — no backend ' + "does today. Without it, (path, size) is all there is to compare against a fresh listing: enough " + "to catch files added, removed, renamed or resized, but not a same-size in-place edit." ), ) num_rows: int | None = Field( @@ -398,13 +401,21 @@ class SamplingInfo(BaseModel): class DatasetProfile(BaseModel): - """The machine-owned dataset profile — the root of the stored contract.""" + """The machine-owned dataset profile — the root of the stored contract. + + Deliberately carries no staleness marker. A stored digest would freeze "which files count as + inputs" into the data at write time, and that judgment moves: once card front-matter drives + split declaration, ``README.md`` becomes an input. Changing the rule would then invalidate every + stored profile at once, with no way to tell a real change from a definition change. The + ``FileRecord``s already describe the inputs, so a consumer that needs to know whether a profile + is current compares them against a fresh listing — same cost, and it learns *what* changed + rather than merely *that* something did. + """ profile_schema_version: str = Field( default=PROFILE_SCHEMA_VERSION, description='Semver of THIS contract (e.g. "1.0") — gates consumer compatibility.', ) - content_digest: str = Field(description="Digest over the stored FileRecords; staleness = mismatch.") created_at: datetime profiler_info: dict = Field( default_factory=dict, diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index ca7ba48484..802cfb988a 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -33,7 +33,6 @@ # --- Fixture: trl-lib/OpenMathReasoning (conversational prompt_completion, verifiable) --- OPENMATHREASONING = """ profile_schema_version: "1.0" -content_digest: sha256:7be1... created_at: 2026-07-08T22:05:12Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 2112, @@ -76,7 +75,6 @@ # --- Fixture: trl-lib/hh-rlhf-helpful-base (conversational preference_pair, explicit) ----- HH_RLHF_HELPFUL_BASE = """ profile_schema_version: "1.0" -content_digest: sha256:5d20... created_at: 2026-07-08T22:41:37Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 1024, @@ -118,7 +116,6 @@ # --- Fixture: nvidia/HelpSteer2 (standard scored_response, no verifiability) -------------- HELPSTEER2 = """ profile_schema_version: "1.0" -content_digest: sha256:c41f... created_at: 2026-07-09T10:12:45Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 1024, @@ -172,7 +169,6 @@ def _build_profile() -> DatasetProfile: """A hand-built profile exercising every model in the contract.""" return DatasetProfile( - content_digest="sha256:deadbeef", created_at=datetime(2026, 7, 13, 12, 0, 0), profiler_info={"name": "nemo-dataset-profiler", "version": "0.1.0"}, sampling=SamplingInfo( @@ -342,6 +338,18 @@ def test_unknown_fields_are_ignored_for_forward_compat(): assert profile.partitions[0].classification.dataset_type == "scored_response" +def test_a_profile_written_before_the_digest_was_dropped_still_loads(): + # `content_digest` was removed rather than repaired: it froze "which files count as inputs" into + # stored data at write time, and that judgment moves. Profiles already written with it have to + # keep loading, or removing it would break every one of them at once — the very failure mode the + # removal exists to avoid. + doc = yaml.safe_load(HELPSTEER2) + doc["content_digest"] = "sha256:7be1c0ffee" + profile = DatasetProfile.model_validate(doc) + assert not hasattr(profile, "content_digest") + assert profile.partitions[0].classification.dataset_type == "scored_response" + + def test_quantiles_and_message_stats_construct(): """Smoke-check the leaf stat models are wired as documented.""" stats = MessageStats( diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py deleted file mode 100644 index 1f9a233795..0000000000 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/digest.py +++ /dev/null @@ -1,28 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""The content digest — a stable fingerprint of a dataset's file listing.""" - -from __future__ import annotations - -import hashlib - -from nemo_datasets_plugin.profiler.file_source import FileEntry - - -def content_digest(entries: list[FileEntry]) -> str: - """A stable digest over the (path, size, checksum) of every file, sorted by path. - - Uses only listing metadata — no file reads — so it is cheap and a profile can self-describe the - inputs it was built from; a mismatch on re-listing means the profile is stale. When a source - reports no checksum, (path, size) is the fallback, which cannot detect a same-size in-place edit. - """ - hasher = hashlib.sha256() - for entry in sorted(entries, key=lambda entry: entry.path): - hasher.update(entry.path.encode("utf-8")) - hasher.update(b"\0") - hasher.update(str(entry.size_bytes).encode("utf-8")) - hasher.update(b"\0") - hasher.update((entry.checksum or "").encode("utf-8")) - hasher.update(b"\n") - return f"sha256:{hasher.hexdigest()}" diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 2eef5bf50b..40869375d7 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -22,7 +22,6 @@ import pyarrow as pa from nemo_datasets_plugin.profiler.classify import classify -from nemo_datasets_plugin.profiler.digest import content_digest from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data @@ -126,8 +125,6 @@ def profile( seed=None, # head sampling makes no random choices; a seed would be theatre ) return DatasetProfile( - # Digest only the files stored as FileRecords, so the profile can recompute its own digest. - content_digest=content_digest(data_entries), created_at=created_at, profiler_info=profiler_info, sampling=sampling, diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index e5f52f0a44..403617dfa7 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -1,14 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for the profiling pipeline: digest, split/partition resolution, and envelope assembly.""" +"""Tests for the profiling pipeline: split/partition resolution and envelope assembly.""" import json from datetime import datetime, timezone import pyarrow as pa import pyarrow.parquet as pq -from nemo_datasets_plugin.profiler.digest import content_digest from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.pipeline import _measure, profile @@ -27,22 +26,6 @@ def _entries(*paths): return [FileEntry(path=p, size_bytes=100) for p in paths] -# --- content digest ------------------------------------------------------------------------------ - - -def test_content_digest_is_stable_and_order_independent(): - a = _entries("train.parquet", "test.parquet") - b = list(reversed(a)) - assert content_digest(a) == content_digest(b) - assert content_digest(a).startswith("sha256:") - - -def test_content_digest_changes_with_size(): - base = _entries("train.parquet") - bigger = [FileEntry(path="train.parquet", size_bytes=200)] - assert content_digest(base) != content_digest(bigger) - - # --- split resolution ---------------------------------------------------------------------------- @@ -128,7 +111,6 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert result.content_digest.startswith("sha256:") assert result.profiler_info["name"] == "nemo-dataset-profiler" assert len(result.partitions) == 1 partition = result.partitions[0] @@ -468,23 +450,23 @@ def test_profile_survives_a_hostile_directory(tmp_path): assert DatasetProfile.model_validate_json(result.model_dump_json()) == result -def test_profile_digest_covers_only_stored_files(tmp_path): - _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) - without_card = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - - (tmp_path / "README.md").write_text("a dataset card") # a non-data file, never stored as a FileRecord - with_card = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - - # A file the profile does not store must not move the digest... - assert without_card.content_digest == with_card.content_digest - # ...and the digest is recomputable from exactly the FileRecords the profile stores. - stored = [ - FileEntry(path=f.path, size_bytes=f.size_bytes, checksum=f.checksum) - for partition in with_card.partitions - for split in partition.splits - for f in split.files - ] - assert content_digest(stored) == with_card.content_digest +def test_stored_file_records_reproduce_the_input_list(tmp_path): + # The contract promises split membership is exhaustive and disjoint, which is what lets a + # consumer compare a stored profile against a fresh listing to decide whether it is current. + # That comparison is the whole reason the records carry path/size/checksum, so the invariant is + # worth asserting directly rather than through a digest that happened to depend on it. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"a": 1}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"a": 2}]) + _write_parquet(tmp_path / "test-00000-of-00001.parquet", [{"a": 3}]) + (tmp_path / "README.md").write_text("a dataset card") # not data; never becomes a FileRecord + + source = LocalFileSource(tmp_path) + result = profile(source, created_at=FIXED_TIME) + + stored = [f.path for partition in result.partitions for split in partition.splits for f in split.files] + listed = [e.path for e in source.list_files() if e.path.endswith(".parquet")] + assert sorted(stored) == sorted(listed) # exhaustive + assert len(stored) == len(set(stored)) # and disjoint def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): From 8c78dde261fbd6ce2ad228cce7a6512f2bd1ade0 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 5 Aug 2026 10:38:33 -0400 Subject: [PATCH 21/44] refactor(datasets): make the directory a partition's identity, not its name A partition's name was doing two jobs -- display label and identity -- and failed at both. Dropping an unrelated notes.jsonl into main/ turned partition "main" into "main:parquet": not renamed, *gone*, so a stored reference resolved to nothing and naming went inconsistent within one profile (socratic bare, main:parquet qualified). Root-level files and a directory literally named "default" collided under the same label with no way to tell them apart. Both trace to PartitionProfile.file_format being a scalar. One format per partition forces a mixed directory to split, splitting forces name qualification to keep the halves distinct, and qualification makes the name depend on which files happen to sit alongside. But format is a property of a *file*; the contract modeled it a layer up, and everything downstream contorted to make that true. So format moves to FileRecord, where it sits beside num_rows and error as another fact about how that file read. The partition summarizes what its files turned out to be, in file_formats, guarded by a subset validator so the summary cannot claim homogeneity a record contradicts. _split_by_format is gone and the reader is resolved per file. A stray .jsonl is noise, not a second dataset: it stays in its directory's partition, contributes its rows, and shows up in the summary. group_partitions now returns the directory rather than a label -- None for root-level files -- and PartitionProfile carries it as source_dir. That is the identity. A lone group under data/ still labels as "default" while its identity stays "data", so the label can move when the layout does without the reference moving with it. name is documented as a label, not a key, and is still not unique: root files and a default/ directory both label "default" and are told apart by source_dir alone. This is where the schema fallback from the previous commit earns its keep. Mixed partitions are now reachable, so a group where only some files declare a schema infers from rows instead -- which is what keeps a column that only the schemaless file witnesses from vanishing, the defect _split_by_format was working around. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 58 ++++++++++++-- .../tests/files/test_dataset_profile.py | 52 ++++++++++--- .../profiler/partition.py | 25 +++--- .../nemo_datasets_plugin/profiler/pipeline.py | 54 +++++++------ plugins/nemo-datasets/tests/test_pipeline.py | 76 +++++++++++++------ 5 files changed, 190 insertions(+), 75 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 43e676a5bd..1bb49d89b5 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -265,6 +265,14 @@ class FileRecord(BaseModel): "to catch files added, removed, renamed or resized, but not a same-size in-place edit." ), ) + file_format: str | None = Field( + default=None, + description=( + "The format this file was read as (jsonl | parquet). A property of the file, not of the " + "partition holding it — which is why a partition may hold more than one. None only on a " + "profile written before formats were recorded per file." + ), + ) num_rows: int | None = Field( default=None, description="Exact only (parquet footer / exhaustive scan), else None.", @@ -322,19 +330,41 @@ class SplitProfile(BaseModel): class PartitionProfile(BaseModel): - """A file-group sharing one row schema, one top-level directory, and one split-name variant - (roughly an HF config); named after the directory / variant, else "default". + """A file-group sharing one row schema and one source directory (roughly an HF config). File membership and row counts live on ``splits`` — every file lands in exactly one split, so partition-level files / num_examples would be derivable duplication. + + ``source_dir`` is the identity and ``name`` is a label. They were a single string until the + consequences showed: root-level files and a directory literally named ``default`` collided under + one label, and dropping an unrelated file into a directory renamed that partition out of + existence — not changed, *gone*, so a stored reference resolved to nothing. """ - name: str = "default" - file_format: str = Field( + name: str = Field( + default="default", + description=( + "Display label, NOT a key. Derived from the layout, not guaranteed unique, and free to " + "change when the layout does. Reference a partition by `source_dir` — or, once card " + "front-matter is parsed, by its declared config name." + ), + ) + source_dir: str | None = Field( + default=None, description=( - "jsonl | parquet are read today; csv | arrow are reserved vocabulary the profiler cannot " - "read yet — files in those formats are reported as unsupported rather than profiled, so " - "this value never appears without a real partition behind it." + "Top-level directory whose files make up this partition; None when they sit at the " + "fileset root. The partition's identity: None and a directory named 'default' are " + 'different partitions even though both label as "default".' + ), + ) + file_formats: list[str] = Field( + default_factory=list, + description=( + "The distinct formats among this partition's files, sorted — normally exactly one. " + "Format is a property of a file (see `FileRecord.file_format`), never a partition " + "dimension: a stray .jsonl beside .parquet shards is noise, not a second dataset, so it " + "stays in this partition and shows up here. jsonl | parquet are read today; csv | arrow " + "are reserved vocabulary the profiler cannot read yet and reports as unsupported." ), ) splits: list[SplitProfile] = Field(description="card-declared > path-detected > single 'default' split.") @@ -359,6 +389,20 @@ def _stats_keys_subset_of_features(self) -> PartitionProfile: raise ValueError(f"stats keys must name top-level features; unknown columns: {sorted(unknown)}") return self + @model_validator(mode="after") + def _file_formats_cover_the_records(self) -> PartitionProfile: + """Every format recorded on a file must appear in the partition's summary. + + A subset check rather than equality, so the two cannot drift in the direction that matters: + a summary omitting a format that demonstrably exists is wrong, while a profile written + before formats were recorded per file has nothing on its records and nothing to check. + """ + recorded = {file.file_format for split in self.splits for file in split.files if file.file_format} + missing = recorded - set(self.file_formats) + if missing: + raise ValueError(f"file_formats omits formats present on this partition's files: {sorted(missing)}") + return self + # ---- envelope ---------------------------------------------------------------------------------- diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index 802cfb988a..f64479708a 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -40,14 +40,15 @@ files_scanned: 33, per_file_row_cap: 64} partitions: - name: default - file_format: parquet + source_dir: null + file_formats: [parquet] splits: - {name: train, canonical: train, num_examples: 3200861, files: [{path: train-00000-of-00032.parquet, size_bytes: 193777041, - checksum: sha256:9c1e..., num_rows: 100027}]} + checksum: sha256:9c1e..., num_rows: 100027, file_format: parquet}]} - {name: test, canonical: test, num_examples: 200, files: [{path: test-00000-of-00001.parquet, size_bytes: 411552, - checksum: sha256:02af..., num_rows: 200}]} + checksum: sha256:02af..., num_rows: 200, file_format: parquet}]} features: - {name: prompt, dtype: messages, semantic_role: prompt, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -81,14 +82,15 @@ rows_total: 46189, files_scanned: 2, per_file_row_cap: 512} partitions: - name: default - file_format: parquet + source_dir: null + file_formats: [parquet] splits: - {name: train, canonical: train, num_examples: 43835, files: [{path: train-00000-of-00001.parquet, size_bytes: 22105331, - checksum: sha256:77b0..., num_rows: 43835}]} + checksum: sha256:77b0..., num_rows: 43835, file_format: parquet}]} - {name: test, canonical: test, num_examples: 2354, files: [{path: test-00000-of-00001.parquet, size_bytes: 1198422, - checksum: sha256:5c1d..., num_rows: 2354}]} + checksum: sha256:5c1d..., num_rows: 2354, file_format: parquet}]} features: - {name: prompt, dtype: messages, semantic_role: prompt, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -122,14 +124,15 @@ rows_total: 21362, files_scanned: 2, per_file_row_cap: 512} partitions: - name: default - file_format: parquet + source_dir: null + file_formats: [parquet] splits: - {name: train, canonical: train, num_examples: 20324, files: [{path: train-00000-of-00001.parquet, size_bytes: 44201991, - checksum: sha256:e410..., num_rows: 20324}]} + checksum: sha256:e410..., num_rows: 20324, file_format: parquet}]} - {name: validation, canonical: validation, num_examples: 1038, files: [{path: validation-00000-of-00001.parquet, size_bytes: 2311008, - checksum: sha256:8bd2..., num_rows: 1038}]} + checksum: sha256:8bd2..., num_rows: 1038, file_format: parquet}]} features: - {name: prompt, dtype: string, semantic_role: prompt} - {name: response, dtype: string, semantic_role: completion} @@ -182,14 +185,21 @@ def _build_profile() -> DatasetProfile: ), partitions=[ PartitionProfile( - file_format="parquet", + source_dir=None, + file_formats=["parquet"], splits=[ SplitProfile( name="train", canonical="train", num_examples=2048, files=[ - FileRecord(path="train-00000.parquet", size_bytes=123, checksum="sha256:ab", num_rows=2048) + FileRecord( + path="train-00000.parquet", + size_bytes=123, + checksum="sha256:ab", + file_format="parquet", + num_rows=2048, + ) ], ) ], @@ -338,6 +348,26 @@ def test_unknown_fields_are_ignored_for_forward_compat(): assert profile.partitions[0].classification.dataset_type == "scored_response" +def test_file_formats_must_not_omit_a_format_its_files_carry(): + # The partition summary is derived from the records, so the two can drift. A summary claiming + # one format while a file says otherwise would report the partition as homogeneous when it is + # not -- the very assumption that made format a partition dimension and cost names their + # stability. Subset, not equality, so a profile written before per-file formats still loads. + doc = yaml.safe_load(HELPSTEER2) + doc["partitions"][0]["splits"][0]["files"][0]["file_format"] = "jsonl" + with pytest.raises(ValueError, match="file_formats omits"): + DatasetProfile.model_validate(doc) + + +def test_a_partition_written_before_per_file_formats_still_loads(): + doc = yaml.safe_load(HELPSTEER2) + for split in doc["partitions"][0]["splits"]: + for file in split["files"]: + del file["file_format"] + profile = DatasetProfile.model_validate(doc) + assert profile.partitions[0].splits[0].files[0].file_format is None + + def test_a_profile_written_before_the_digest_was_dropped_still_loads(): # `content_digest` was removed rather than repaired: it froze "which files count as inputs" into # stored data at write time, and that judgment moves. Profiles already written with it have to diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py index 33c7c84b97..0703c46cef 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py @@ -29,20 +29,25 @@ def _top_dir(path: str) -> str | None: return parts[0] -def group_partitions(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry]]]: - """Group files into (name, files) partitions by top-level directory. - - A single top-level group — every file at the root, or all under one container like ``data/`` — - is one "default" partition. Multiple top-level directories (e.g. ``main/`` and ``socratic/``) - each become their own partition, named after the directory; root-level files fall under - "default". +def group_partitions(entries: list[FileEntry]) -> list[tuple[str | None, list[FileEntry]]]: + """Group files into (source_dir, files) partitions by top-level directory. + + Returns the *directory*, not a label — ``None`` for files at the fileset root. The label is the + caller's to derive, because the two are not the same thing: a lone group under ``data/`` labels + as "default" while its identity stays ``"data"``, and root-level files are a different partition + from a directory literally named ``default`` even though both label the same way. Collapsing + those into one string is what let a partition collide with another, and what let an unrelated + file rename one out of existence. + + Files whose top-level directory is a split name (``train/``, ``test/``) group under ``None`` + alongside root-level files: those are one dataset's splits, not separate partitions. """ by_dir: dict[str | None, list[FileEntry]] = {} for entry in entries: by_dir.setdefault(_top_dir(entry.path), []).append(entry) if len(by_dir) == 1: - return [("default", list(entries))] + # A single group is one partition holding everything, whatever its directory happened to be. + return [(next(iter(by_dir)), list(entries))] - ordered = sorted(by_dir.items(), key=lambda item: (item[0] is None, item[0] or "")) - return [("default" if directory is None else directory, files) for directory, files in ordered] + return sorted(by_dir.items(), key=lambda item: (item[0] is None, item[0] or "")) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 40869375d7..e389496b9c 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -91,17 +91,14 @@ def profile( files_read = 0 all_scanned = True - for partition_name, partition_entries in group_partitions(data_entries): - format_groups = _split_by_format(partition_entries) - for file_format, format_entries in format_groups: - # A directory that holds more than one format yields one partition per format; qualify - # the name so the partitions stay distinct. A single-format directory keeps its bare name. - name = f"{partition_name}:{file_format}" if len(format_groups) > 1 else partition_name - outcome = _profile_partition(source, name, file_format, format_entries, row_cap) - partitions.append(outcome.partition) - rows_scanned += outcome.rows_scanned - files_read += outcome.files_read - all_scanned = all_scanned and outcome.scanned_all + groups = group_partitions(data_entries) + for source_dir, partition_entries in groups: + name = _partition_label(source_dir, len(groups)) + outcome = _profile_partition(source, name, source_dir, partition_entries, row_cap) + partitions.append(outcome.partition) + rows_scanned += outcome.rows_scanned + files_read += outcome.files_read + all_scanned = all_scanned and outcome.scanned_all # Data we could not read is data we did not scan, so unsupported files defeat exhaustiveness just # as an unreadable file does. @@ -132,18 +129,16 @@ def profile( ) -def _split_by_format(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry]]]: - """Sub-group a directory's files by format so each profiled partition is format-homogeneous. +def _partition_label(source_dir: str | None, group_count: int) -> str: + """The display label for a partition. Cosmetic only — ``source_dir`` carries the identity. - ``group_partitions`` groups by directory only, but one directory can hold more than one format - (a stray ``.jsonl`` beside ``.parquet`` shards). Left mixed, a partition would derive its schema - from whichever format was read first and then measure rows from both. Sorted for deterministic - partition order; ``entries`` are pre-filtered ``data_entries`` so every format is registered. + A lone group labels as "default" whatever directory it came from, so the common + everything-under-``data/`` layout does not surface a meaningless container name. With several + groups each takes its directory name, and root-level files take "default". """ - by_format: dict[str, list[FileEntry]] = {} - for entry in entries: - by_format.setdefault(_format_of(entry.path), []).append(entry) - return sorted(by_format.items()) + if group_count == 1 or source_dir is None: + return "default" + return source_dir def _unify_schemas(schemas: list[pa.Schema]) -> pa.Schema | None: @@ -214,9 +209,15 @@ class _PartitionOutcome: def _profile_partition( - source: FileSource, name: str, file_format: str, entries: list[FileEntry], row_cap: int | None + source: FileSource, name: str, source_dir: str | None, entries: list[FileEntry], row_cap: int | None ) -> _PartitionOutcome: - """Profile one format-homogeneous partition. + """Profile one partition — the files of one source directory, whatever formats they are in. + + The reader is resolved per file rather than per partition. Format is a property of a file, and + a directory holding two of them is a stray file, not a second dataset; splitting the partition + to keep one scalar ``file_format`` true is what made partition names unstable. Mixed formats + instead flow through to ``_measure``, which infers the schema from rows when not every file + declared one. An unreadable file (or a format with no registered reader) is isolated: it keeps its FileRecord, records *why* on ``FileRecord.error``, contributes no rows, and flips ``scanned_all`` off — it @@ -235,6 +236,7 @@ def _profile_partition( split_counts_known = True # every file's exact total row count is known (footer or full scan) split_scanned = True # every row of every file was actually parsed for entry in split.entries: + file_format = _format_of(entry.path) error: str | None = None try: result = get_reader(file_format).read(source, entry, row_cap=row_cap) @@ -267,6 +269,7 @@ def _profile_partition( path=entry.path, size_bytes=entry.size_bytes, checksum=entry.checksum, + file_format=file_format, num_rows=num_rows, error=error, ) @@ -291,7 +294,10 @@ def _profile_partition( ) partition = PartitionProfile( name=name, - file_format=file_format, + source_dir=source_dir, + # Summarized from the records rather than assumed: the partition no longer picks a format, + # it reports the ones its files turned out to be in. + file_formats=sorted({file.file_format for split in split_profiles for file in split.files if file.file_format}), splits=split_profiles, features=features, stats=stats, diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 403617dfa7..b6bdfd8541 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -68,24 +68,26 @@ def test_resolve_splits_falls_back_to_single_default(): def test_group_partitions_single_default_for_root_files(): assert group_partitions(_entries("train.parquet", "test.parquet")) == [ - ("default", _entries("train.parquet", "test.parquet")) + (None, _entries("train.parquet", "test.parquet")) ] def test_group_partitions_collapses_single_container_dir(): + # One container directory is still one partition, but its identity stays the directory. Losing + # "data" here is what let a partition's identity move when the surrounding layout changed. parts = group_partitions(_entries("data/train.parquet", "data/test.parquet")) - assert [name for name, _ in parts] == ["default"] + assert [source_dir for source_dir, _ in parts] == ["data"] def test_group_partitions_splits_multiple_top_dirs(): parts = group_partitions(_entries("main/train.parquet", "socratic/train.parquet")) - assert [name for name, _ in parts] == ["main", "socratic"] + assert [source_dir for source_dir, _ in parts] == ["main", "socratic"] def test_group_partitions_does_not_treat_split_dirs_as_partitions(): # train/ and test/ are one dataset's splits, not two datasets. parts = group_partitions(_entries("train/data.parquet", "test/data.parquet")) - assert [name for name, _ in parts] == ["default"] + assert [source_dir for source_dir, _ in parts] == [None] def test_resolve_splits_reads_the_split_directory(): @@ -115,7 +117,7 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert len(result.partitions) == 1 partition = result.partitions[0] assert partition.name == "default" - assert partition.file_format == "parquet" + assert partition.file_formats == ["parquet"] splits = {s.name: s for s in partition.splits} assert set(splits) == {"train", "validation"} @@ -146,7 +148,7 @@ def test_profile_jsonl_dataset_counts_rows_exactly(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) partition = result.partitions[0] - assert partition.file_format == "jsonl" + assert partition.file_formats == ["jsonl"] assert partition.splits[0].name == "train" assert partition.splits[0].num_examples == 3 assert result.sampling.rows_scanned == 3 @@ -159,7 +161,8 @@ def test_profile_multiple_directories_become_partitions(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) assert [p.name for p in result.partitions] == ["main", "socratic"] - assert all(p.file_format == "parquet" for p in result.partitions) + assert all(p.file_formats == ["parquet"] for p in result.partitions) + assert [p.source_dir for p in result.partitions] == ["main", "socratic"] def test_profile_top_level_split_dirs_become_one_partition(tmp_path): @@ -192,26 +195,51 @@ def test_profile_nested_split_dirs_keep_splits_apart(tmp_path): assert splits["test"].num_examples == 1 -def test_profile_splits_mixed_formats_into_separate_partitions(tmp_path): - # A directory holding two formats must not profile as one partition: features/stats would be - # derived from one format's schema but measured over rows from both. Each format becomes its own - # format-homogeneous partition, name-qualified so the two stay distinct. +def test_profile_keeps_a_mixed_format_directory_as_one_partition(tmp_path): + # A stray .jsonl beside .parquet shards is noise, not a second dataset. Splitting the partition + # to keep a scalar `file_format` true invented structure that is not in the data *and* renamed + # the real partition (default -> default:parquet). Format is a per-file fact instead. _write_parquet(tmp_path / "data" / "train-00000-of-00001.parquet", [{"prompt": "a"}]) (tmp_path / "data" / "extra.jsonl").write_text('{"question": "b"}\n{"question": "c"}\n') result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - by_format = {p.file_format: p for p in result.partitions} - assert set(by_format) == {"parquet", "jsonl"} - assert by_format["parquet"].name == "default:parquet" - assert by_format["jsonl"].name == "default:jsonl" - # Each partition's schema reflects only its own files. - assert [f.name for f in by_format["parquet"].features] == ["prompt"] - assert [f.name for f in by_format["jsonl"].features] == ["question"] + assert len(result.partitions) == 1 + partition = result.partitions[0] + assert partition.name == "default" + assert partition.source_dir == "data" + assert partition.file_formats == ["jsonl", "parquet"] + assert {f.path.rsplit("/", 1)[-1]: f.file_format for s in partition.splits for f in s.files} == { + "train-00000-of-00001.parquet": "parquet", + "extra.jsonl": "jsonl", + } + # Both formats' columns reach features. Trusting the declared parquet schema would have erased + # `question`, which only the schemaless file witnesses -- the defect the split worked around. + assert sorted(f.name for f in partition.features) == ["prompt", "question"] assert result.sampling.rows_scanned == 3 # 1 parquet + 2 jsonl, each counted once assert result.sampling.exhaustive is True +def test_root_files_and_a_directory_named_default_stay_distinct(): + # Both label as "default"; only source_dir tells them apart. Flattening the two into one string + # produced two partitions with the same name and no way to reference either. + parts = group_partitions(_entries("root.parquet", "default/inner.parquet")) + assert [source_dir for source_dir, _ in parts] == ["default", None] + + +def test_an_unrelated_file_does_not_rename_a_partition(tmp_path): + # Dropping a stray .jsonl into main/ used to turn partition "main" into "main:parquet" -- not + # renamed, *gone*, so a stored reference resolved to nothing. + _write_parquet(tmp_path / "main" / "train.parquet", [{"q": "a"}]) + _write_parquet(tmp_path / "socratic" / "train.parquet", [{"q": "b"}]) + before = [(p.name, p.source_dir) for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] + + (tmp_path / "main" / "notes.jsonl").write_text('{"note": "someone dropped this here"}\n') + after = [(p.name, p.source_dir) for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] + + assert before == after == [("main", "main"), ("socratic", "socratic")] + + def test_profile_unions_columns_across_shards(tmp_path): # A column that appears only in a later shard must still reach features/stats. Taking the first # shard's schema would drop it entirely. @@ -439,12 +467,14 @@ def test_profile_survives_a_hostile_directory(tmp_path): assert records["train-00001-of-00002.parquet"].error is not None # corrupt file, named and explained assert records["extra.jsonl"].error is not None # partial parse, named and explained - # The readable parquet rows still produced a real classification. - parquet_partition = next(p for p in result.partitions if p.file_format == "parquet") - assert parquet_partition.classification.dataset_type == "prompt_completion" + # One partition, not one per format: the stray .jsonl is noise, not a second dataset. + assert len(result.partitions) == 1 + partition = result.partitions[0] + assert partition.file_formats == ["jsonl", "parquet"] + # The readable parquet rows still produced a real classification... + assert partition.classification.dataset_type == "prompt_completion" # ...and the odd jsonl rows were measured rather than aborting the run. - jsonl_partition = next(p for p in result.partitions if p.file_format == "jsonl") - assert jsonl_partition.stats["messages"].messages.roles_seen == ["1", "user"] + assert partition.stats["messages"].messages.roles_seen == ["1", "user"] # The whole thing still round-trips as a stored profile. assert DatasetProfile.model_validate_json(result.model_dump_json()) == result From 046ecb2add34a45de0128aa7988a1a8e849c561f Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 5 Aug 2026 11:16:13 -0400 Subject: [PATCH 22/44] refactor(datasets): split coverage into the numbers it was collapsing `SamplingInfo.exhaustive` was one bit answering two questions. "Are these measurements facts or estimates?" is a property of a partition's stats; "did I see all the data?" is a property of the fileset's inputs. It was stored where neither was decided, and it folded together causes that call for different people to act -- a row cap is the caller's choice, a corrupt shard the data owner's, a missing reader ours. It also did not describe what it appeared to. The value gating whether categorical.values may quote a proven enumeration was `partition_scanned`, computed per partition and never stored, so one corrupt file in socratic/ set the envelope to false while main/ went on quoting its enumeration -- correctly. So the bit becomes PartitionProfile.stats_complete, scoped to the measurements it qualifies and equal to the value that was driving the decision all along. The fileset question becomes counts with denominators: rows_scanned/rows_present and files_read/files_present. `files_scanned` previously told consumers to "expect this to equal the fileset's file count" -- a number the profile did not carry. rows_present is gated on the count being *known*, not on the read being complete. The old rule nulled it exactly when it carried information: a capped parquet run knows its totals from the footers, so "4 of 10" was reportable and came out as "4 of unknown", while a complete scan made it equal rows_scanned and say nothing. SplitProfile.num_examples had the mirror-image problem in prose, warning about extrapolation the implementation never does. Unsupported files stop being bare paths in the free-form profiler_info dict and become real FileRecords on DatasetProfile.unreadable_files, each carrying its reason on `error` -- the same representation as any other file the profiler could not read. They differ only in living at the envelope, because no partition ever grouped them. `strategy` follows format down to FileRecord.read_strategy: the cap is applied per file, and head-vs-full is a per-file read decision now that a partition can hold more than one format. Its documented vocabulary was wrong anyway, listing stratified_probes | random while the profiler emitted head_per_file. The dataset-wide question is still one expression, and now names which half failed: all(p.stats_complete for p in profile.partitions) and not profile.unreadable_files PROFILE_SCHEMA_VERSION stays at 1.0. The fields have moved a great deal across this and the two preceding commits, but nothing consumes the contract yet, so there is no compatibility to gate; the first version number that means anything is the one shipped alongside the first consumer. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 90 ++++++++++++----- .../tests/files/test_dataset_profile.py | 37 +++---- .../nemo_datasets_plugin/profiler/pipeline.py | 74 +++++++++----- plugins/nemo-datasets/tests/test_pipeline.py | 97 +++++++++++++------ 4 files changed, 203 insertions(+), 95 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 1bb49d89b5..059cd3bc25 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -30,7 +30,9 @@ from pydantic import BaseModel, Field, model_validator # Semver of THIS contract. Gates consumer compatibility: new detectors or vocabulary values are a -# minor bump; a change to the fields below is a major bump. +# minor bump; a change to the fields below is a major bump. Still 1.0 because nothing consumes it +# yet — the fields have moved a great deal, but pre-release churn is not a break for anyone, and the +# first number that means something is the one shipped alongside the first consumer. PROFILE_SCHEMA_VERSION = "1.0" @@ -273,6 +275,16 @@ class FileRecord(BaseModel): "profile written before formats were recorded per file." ), ) + read_strategy: str | None = Field( + default=None, + description=( + "How this file's rows were sampled: full | head. The *policy* applied, not the outcome — " + "a head-capped read of a file smaller than the cap still says head, and whether it ended " + "up complete is `num_rows` versus what was scanned. Per file because it follows format, " + "which is also per file: a parquet shard can be sampled by row group where a jsonl file " + "in the same partition can only be read from the top." + ), + ) num_rows: int | None = Field( default=None, description="Exact only (parquet footer / exhaustive scan), else None.", @@ -322,9 +334,10 @@ class SplitProfile(BaseModel): num_examples: int | None = Field( default=None, description=( - "Rows in this split, counting every file in `files` whether or not it was scanned. Exact when " - "read from parquet footers or an exhaustive scan, otherwise extrapolated from the rows sampled — " - "check `SamplingInfo.exhaustive` before treating it as a fact. None when nothing usable was found." + "Rows in this split, counting every file in `files` whether or not its rows were read. Always " + "exact — summed from parquet footers or from files read to their end — and None the moment any " + "one file's count is unknown. Never an estimate, so it carries no accuracy caveat: a capped run " + "still reports the true total whenever the footers knew it." ), ) @@ -378,6 +391,16 @@ class PartitionProfile(BaseModel): "omitted); keys are a subset of the top-level `features` names." ), ) + stats_complete: bool = Field( + description=( + "True => `features`, `stats` and `classification` were computed over every row of every " + "file in THIS partition: proven facts, not estimates. Only then can a consumer assert " + "enum / required in a bridged JSON Schema, or read a verifiability coverage of 1.0 as " + "literal. Scoped to the partition because that is where it is decided — a corrupt shard " + "in one partition says nothing about the measurements in another, and a fileset-wide " + "flag quietly downgraded every partition to the worst one." + ), + ) classification: PartitionClassification @model_validator(mode="after") @@ -408,36 +431,42 @@ def _file_formats_cover_the_records(self) -> PartitionProfile: class SamplingInfo(BaseModel): - """How much of the data the profile is based on. + """How much of the data the profile is based on — coverage, stated as numbers. + + Deliberately carries no ``exhaustive`` flag. That bit was answering two questions at once: "are + these measurements facts or estimates?", which is a property of a *partition* and now lives on + ``PartitionProfile.stats_complete``, and "did I see all the data?", which is this block's job and + needs numerators and denominators rather than a boolean. It also folded together causes that + call for different people to act — a row cap is the caller's choice, a corrupt shard is the data + owner's problem, and a missing reader is ours. + + The dataset-wide question is still one expression away, and now says which half failed:: - Consumers read ``exhaustive`` to decide whether stats are proven facts or estimates (e.g. only an - exhaustive profile can assert enum / required in a bridged JSON Schema, or that verifiability - coverage is truly 1.0). + all(p.stats_complete for p in profile.partitions) and not profile.unreadable_files """ - exhaustive: bool = Field(description="True => every row of every file was parsed.") - strategy: str = Field( - description=( - "full | stratified_probes | random. Kept explicit alongside `exhaustive` because, with an open " - "strategy vocabulary, consumers can't derive exhaustiveness from the strategy name alone." - ), - ) rows_scanned: int = Field(description="Total rows actually parsed across all files.") - rows_total: int | None = Field( + rows_present: int | None = Field( default=None, description=( - "How many rows the whole fileset holds, scanned or not — the denominator `rows_scanned` is a " - "fraction of, so a consumer can judge how representative the stats are. Populated only when the " - "count is exact and cheap (summed parquet footers, or an exhaustive scan); None means unknown, " - "never zero and never an estimate." + "How many rows the fileset holds, scanned or not — the denominator `rows_scanned` is a " + "fraction of. Populated whenever every file's count is *known*, regardless of how much was " + "read: a row-capped run over parquet still knows its totals from the footers, and that is " + "exactly when the ratio carries information. None means at least one file's count is " + "unknown — never zero, never an estimate." ), ) - files_scanned: int = Field( + files_read: int = Field( + description="Files actually opened and read from (a count; the files themselves are `SplitProfile.files`)." + ) + files_present: int = Field( description=( - "How many files were opened and read from (a count, not a list — the files themselves are " - "`SplitProfile.files`). Every file should be probed, since head-sampling a subset hides columns " - "that appear only in later shards; expect this to equal the fileset's file count, and be lower " - "only when scale forces file-level sampling." + "Data files the fileset holds, whether or not this run could read them — the denominator " + "`files_read` is a fraction of. Includes files in formats with no reader, since those are " + "data that went unprofiled (they are listed in `DatasetProfile.unreadable_files`). A README " + "is not data and is counted nowhere. Every readable file should be opened, since " + "head-sampling a *subset of files* hides columns that appear only in later shards, so expect " + "these two to match until scale forces file-level sampling." ), ) per_file_row_cap: int | None = Field(default=None, description="Cap that bounded per-file reads, if any.") @@ -469,6 +498,17 @@ class DatasetProfile(BaseModel): partitions: list[PartitionProfile] = Field( description="Single partition in the common homogeneous case; there is no fileset-level rollup.", ) + unreadable_files: list[FileRecord] = Field( + default_factory=list, + description=( + "Files that plainly hold dataset records but that no partition could take, because the " + "profiler has no reader for their format; each carries the reason on `error`. Reporting " + "them is what keeps a directory of .csv shards from profiling as an exhaustively scanned " + "*empty* dataset, indistinguishable from one that really is empty. A file whose format is " + "known but whose read failed keeps its FileRecord inside its split instead — it was " + "grouped and attempted, these never were." + ), + ) # Resolve the recursive FeatureSchema self-reference (deferred by `from __future__ import annotations`). diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index f64479708a..ebcd4ebc10 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -35,20 +35,20 @@ profile_schema_version: "1.0" created_at: 2026-07-08T22:05:12Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} -sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 2112, - rows_total: 3201061, - files_scanned: 33, per_file_row_cap: 64} +sampling: {rows_scanned: 2112, rows_present: 3201061, + files_read: 33, files_present: 33, per_file_row_cap: 64} partitions: - name: default source_dir: null file_formats: [parquet] + stats_complete: false splits: - {name: train, canonical: train, num_examples: 3200861, files: [{path: train-00000-of-00032.parquet, size_bytes: 193777041, - checksum: sha256:9c1e..., num_rows: 100027, file_format: parquet}]} + checksum: sha256:9c1e..., num_rows: 100027, file_format: parquet, read_strategy: head}]} - {name: test, canonical: test, num_examples: 200, files: [{path: test-00000-of-00001.parquet, size_bytes: 411552, - checksum: sha256:02af..., num_rows: 200, file_format: parquet}]} + checksum: sha256:02af..., num_rows: 200, file_format: parquet, read_strategy: head}]} features: - {name: prompt, dtype: messages, semantic_role: prompt, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -78,19 +78,20 @@ profile_schema_version: "1.0" created_at: 2026-07-08T22:41:37Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} -sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 1024, - rows_total: 46189, files_scanned: 2, per_file_row_cap: 512} +sampling: {rows_scanned: 1024, rows_present: 46189, + files_read: 2, files_present: 2, per_file_row_cap: 512} partitions: - name: default source_dir: null file_formats: [parquet] + stats_complete: false splits: - {name: train, canonical: train, num_examples: 43835, files: [{path: train-00000-of-00001.parquet, size_bytes: 22105331, - checksum: sha256:77b0..., num_rows: 43835, file_format: parquet}]} + checksum: sha256:77b0..., num_rows: 43835, file_format: parquet, read_strategy: head}]} - {name: test, canonical: test, num_examples: 2354, files: [{path: test-00000-of-00001.parquet, size_bytes: 1198422, - checksum: sha256:5c1d..., num_rows: 2354, file_format: parquet}]} + checksum: sha256:5c1d..., num_rows: 2354, file_format: parquet, read_strategy: head}]} features: - {name: prompt, dtype: messages, semantic_role: prompt, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -120,19 +121,20 @@ profile_schema_version: "1.0" created_at: 2026-07-09T10:12:45Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} -sampling: {exhaustive: false, strategy: stratified_probes, rows_scanned: 1024, - rows_total: 21362, files_scanned: 2, per_file_row_cap: 512} +sampling: {rows_scanned: 1024, rows_present: 21362, + files_read: 2, files_present: 2, per_file_row_cap: 512} partitions: - name: default source_dir: null file_formats: [parquet] + stats_complete: false splits: - {name: train, canonical: train, num_examples: 20324, files: [{path: train-00000-of-00001.parquet, size_bytes: 44201991, - checksum: sha256:e410..., num_rows: 20324, file_format: parquet}]} + checksum: sha256:e410..., num_rows: 20324, file_format: parquet, read_strategy: head}]} - {name: validation, canonical: validation, num_examples: 1038, files: [{path: validation-00000-of-00001.parquet, size_bytes: 2311008, - checksum: sha256:8bd2..., num_rows: 1038, file_format: parquet}]} + checksum: sha256:8bd2..., num_rows: 1038, file_format: parquet, read_strategy: head}]} features: - {name: prompt, dtype: string, semantic_role: prompt} - {name: response, dtype: string, semantic_role: completion} @@ -175,11 +177,10 @@ def _build_profile() -> DatasetProfile: created_at=datetime(2026, 7, 13, 12, 0, 0), profiler_info={"name": "nemo-dataset-profiler", "version": "0.1.0"}, sampling=SamplingInfo( - exhaustive=False, - strategy="stratified_probes", rows_scanned=1024, - rows_total=2048, - files_scanned=2, + rows_present=2048, + files_read=2, + files_present=2, per_file_row_cap=512, seed=7, ), @@ -187,6 +188,7 @@ def _build_profile() -> DatasetProfile: PartitionProfile( source_dir=None, file_formats=["parquet"], + stats_complete=False, splits=[ SplitProfile( name="train", @@ -198,6 +200,7 @@ def _build_profile() -> DatasetProfile: size_bytes=123, checksum="sha256:ab", file_format="parquet", + read_strategy="head", num_rows=2048, ) ], diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index e389496b9c..8213883379 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -19,6 +19,7 @@ from dataclasses import dataclass from datetime import datetime, timezone +from pathlib import PurePosixPath import pyarrow as pa from nemo_datasets_plugin.profiler.classify import classify @@ -81,15 +82,25 @@ def profile( data_entries = [entry for entry in all_entries if detect_format(entry.path) is not None] # Files that plainly hold records but have no reader yet. They are not profiled, but they must be # reported: silently dropping them let a directory of .csv shards profile as an exhaustively - # scanned, empty dataset — indistinguishable from a dataset that really is empty. - unsupported = sorted( - entry.path for entry in all_entries if detect_format(entry.path) is None and is_unsupported_data(entry.path) - ) + # scanned, empty dataset — indistinguishable from a dataset that really is empty. They get real + # FileRecords like any other file the profiler could not read, just at the envelope, since no + # partition ever grouped them. + unreadable_files = [ + FileRecord( + path=entry.path, + size_bytes=entry.size_bytes, + checksum=entry.checksum, + error=f"no reader for '{PurePosixPath(entry.path).suffix.lower()}' files", + ) + for entry in sorted(all_entries, key=lambda entry: entry.path) + if detect_format(entry.path) is None and is_unsupported_data(entry.path) + ] partitions: list[PartitionProfile] = [] rows_scanned = 0 files_read = 0 - all_scanned = True + # None once any file's row count is unknown: the fileset's total is then unknowable, not zero. + rows_present: int | None = 0 groups = group_partitions(data_entries) for source_dir, partition_entries in groups: @@ -98,37 +109,43 @@ def profile( partitions.append(outcome.partition) rows_scanned += outcome.rows_scanned files_read += outcome.files_read - all_scanned = all_scanned and outcome.scanned_all + rows_present = _add_known(rows_present, outcome.rows_present) - # Data we could not read is data we did not scan, so unsupported files defeat exhaustiveness just - # as an unreadable file does. - exhaustive = all_scanned and not unsupported - profiler_info: dict = {"name": PROFILER_NAME, "version": PROFILER_VERSION} - if unsupported: - profiler_info["unsupported_files"] = unsupported + # A format with no reader holds an unknown number of rows, so it makes the fileset total unknown + # in exactly the way an unread file does. + if unreadable_files: + rows_present = None sampling = SamplingInfo( - exhaustive=exhaustive, - # The policy in effect, which `exhaustive` deliberately does not encode: a capped run over - # files that all fit under the cap is still a full scan, and an uncapped run can still fall - # short of exhaustive because a file was unreadable. - strategy="full" if row_cap is None else "head_per_file", - # `rows_total` is documented as never zero: a 0 here would read as "this dataset is empty" - # when it more often means nothing was recognized. Unknown is the honest answer. - rows_total=rows_scanned if exhaustive and rows_scanned else None, rows_scanned=rows_scanned, - files_scanned=files_read, # files actually opened and read, not files merely listed + rows_present=rows_present, + files_read=files_read, # files actually opened and read, not files merely listed + # Every data file, readable or not: the denominator that makes `files_read` a fraction rather + # than a bare count. Non-data files (a README, a LICENSE) are not data and are counted nowhere. + files_present=len(data_entries) + len(unreadable_files), per_file_row_cap=row_cap, seed=None, # head sampling makes no random choices; a seed would be theatre ) return DatasetProfile( created_at=created_at, - profiler_info=profiler_info, + profiler_info={"name": PROFILER_NAME, "version": PROFILER_VERSION}, sampling=sampling, partitions=partitions, + unreadable_files=unreadable_files, ) +def _add_known(total: int | None, addend: int | None) -> int | None: + """Sum two counts, where ``None`` means unknown and poisons the total. + + A fileset whose row count is unknown for even one file has an unknown total — reporting the sum + of the rest would look like a fact and read low. + """ + if total is None or addend is None: + return None + return total + addend + + def _partition_label(source_dir: str | None, group_count: int) -> str: """The display label for a partition. Cosmetic only — ``source_dir`` carries the identity. @@ -204,8 +221,8 @@ class _PartitionOutcome: partition: PartitionProfile rows_scanned: int - files_read: int # files actually opened and read, so `files_scanned` can exclude failures - scanned_all: bool + files_read: int # files actually opened and read, so `files_read` can exclude failures + rows_present: int | None # rows known to exist here, or None once any file's count is unknown def _profile_partition( @@ -228,7 +245,9 @@ def _profile_partition( all_declared = True # every file that contributed rows carried a declared schema rows_scanned = 0 files_read = 0 + rows_present: int | None = 0 partition_scanned = True + read_strategy = "full" if row_cap is None else "head" split_profiles: list[SplitProfile] = [] for split in resolve_splits(entries): file_records: list[FileRecord] = [] @@ -270,10 +289,12 @@ def _profile_partition( size_bytes=entry.size_bytes, checksum=entry.checksum, file_format=file_format, + read_strategy=read_strategy, num_rows=num_rows, error=error, ) ) + rows_present = _add_known(rows_present, num_rows) if num_rows is None: split_counts_known = False else: @@ -301,8 +322,11 @@ def _profile_partition( splits=split_profiles, features=features, stats=stats, + # Scoped to this partition, which is where it was decided all along: `partition_scanned` is + # the value that already gated whether `categorical.values` could quote a proven enumeration. + stats_complete=partition_scanned, classification=classification, ) return _PartitionOutcome( - partition=partition, rows_scanned=rows_scanned, files_read=files_read, scanned_all=partition_scanned + partition=partition, rows_scanned=rows_scanned, files_read=files_read, rows_present=rows_present ) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index b6bdfd8541..fd0d859ed8 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -133,14 +133,15 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert partition.stats["prompt"].text is not None assert partition.classification.dataset_type == "prompt_only" # a lone prompt column, no target - # strategy is the policy, exhaustive is the outcome: a capped run over files that all fit under - # the cap is still a full scan, which is why the contract keeps the two fields independent. - assert result.sampling.strategy == "head_per_file" - assert result.sampling.exhaustive is True + # read_strategy is the policy, stats_complete is the outcome: a capped run over files that all + # fit under the cap is still a complete scan, which is why the two live apart -- and now at the + # levels where each is decided, per file and per partition. + assert {f.read_strategy for s in partition.splits for f in s.files} == {"head"} + assert partition.stats_complete is True assert result.sampling.per_file_row_cap == 1000 assert result.sampling.rows_scanned == 3 - assert result.sampling.rows_total == 3 - assert result.sampling.files_scanned == 2 + assert result.sampling.rows_present == 3 + assert result.sampling.files_read == result.sampling.files_present == 2 def test_profile_jsonl_dataset_counts_rows_exactly(tmp_path): @@ -217,7 +218,7 @@ def test_profile_keeps_a_mixed_format_directory_as_one_partition(tmp_path): # `question`, which only the schemaless file witnesses -- the defect the split worked around. assert sorted(f.name for f in partition.features) == ["prompt", "question"] assert result.sampling.rows_scanned == 3 # 1 parquet + 2 jsonl, each counted once - assert result.sampling.exhaustive is True + assert partition.stats_complete is True def test_root_files_and_a_directory_named_default_stay_distinct(): @@ -297,9 +298,10 @@ def test_profile_isolates_unreadable_files(tmp_path): assert splits["test"].num_examples is None # unreadable -> count unknown, not a crash assert splits["test"].files[0].num_rows is None assert splits["test"].files[0].error is not None # ...and the profile says why - assert result.sampling.exhaustive is False # a file could not be fully parsed - assert result.sampling.rows_total is None - assert result.sampling.files_scanned == 1 # one file was actually read; the other never opened + assert result.partitions[0].stats_complete is False # a file could not be fully parsed + assert result.sampling.rows_present is None + assert result.sampling.files_read == 1 # one file was actually read; the other never opened + assert result.sampling.files_present == 2 # ...out of two that were there to read def test_profile_row_cap_bounds_reads_and_says_so(tmp_path): @@ -309,8 +311,10 @@ def test_profile_row_cap_bounds_reads_and_says_so(tmp_path): assert result.sampling.rows_scanned == 4 assert result.sampling.per_file_row_cap == 4 - assert result.sampling.exhaustive is False # 4 of 10 rows is not a full scan - assert result.sampling.rows_total is None + assert result.partitions[0].stats_complete is False # 4 of 10 rows is not a full scan + # The footer knows the total even though the cap stopped the read. Gating this on completeness + # nulled it exactly when it carried information: "4 of 10" is a ratio, "4 of unknown" is not. + assert result.sampling.rows_present == 10 assert result.partitions[0].splits[0].num_examples == 10 # the footer count survives sampling @@ -319,10 +323,10 @@ def test_profile_uncapped_read_is_a_full_scan(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=None) - assert result.sampling.strategy == "full" + assert {f.read_strategy for s in result.partitions[0].splits for f in s.files} == {"full"} assert result.sampling.per_file_row_cap is None - assert result.sampling.exhaustive is True - assert result.sampling.rows_scanned == 10 + assert result.partitions[0].stats_complete is True + assert result.sampling.rows_scanned == result.sampling.rows_present == 10 def test_profile_cap_larger_than_a_jsonl_file_keeps_it_exhaustive(tmp_path): @@ -333,8 +337,8 @@ def test_profile_cap_larger_than_a_jsonl_file_keeps_it_exhaustive(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=1000) assert result.partitions[0].splits[0].num_examples == 2 - assert result.sampling.exhaustive is True - assert result.sampling.rows_total == 2 + assert result.partitions[0].stats_complete is True + assert result.sampling.rows_present == 2 def test_profile_reports_unsupported_data_files(tmp_path): @@ -346,9 +350,12 @@ def test_profile_reports_unsupported_data_files(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) assert result.partitions == [] - assert result.sampling.exhaustive is False # we scanned nothing, and admit it - assert result.sampling.rows_total is None # not 0: "empty" would be a lie - assert result.profiler_info["unsupported_files"] == ["test.arrow", "train.csv"] + assert result.sampling.rows_present is None # not 0: "empty" would be a lie + assert result.sampling.files_read == 0 + assert result.sampling.files_present == 2 # both are data; neither could be read + # Typed records now, each saying why -- not bare paths tucked into a free-form dict. + assert [f.path for f in result.unreadable_files] == ["test.arrow", "train.csv"] + assert all("no reader" in f.error for f in result.unreadable_files) def test_profile_ignores_non_data_files_without_penalty(tmp_path): @@ -359,8 +366,9 @@ def test_profile_ignores_non_data_files_without_penalty(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert result.sampling.exhaustive is True - assert "unsupported_files" not in result.profiler_info + assert result.partitions[0].stats_complete is True + assert result.unreadable_files == [] + assert result.sampling.files_present == 1 # the README and LICENSE are not data, counted nowhere def test_profile_records_a_partial_jsonl_read(tmp_path): @@ -373,7 +381,7 @@ def test_profile_records_a_partial_jsonl_read(tmp_path): record = result.partitions[0].splits[0].files[0] assert record.num_rows == 2 # the readable rows survived assert record.error is not None and "line 2" in record.error - assert result.sampling.exhaustive is False # a line was lost, so this is not a full scan + assert result.partitions[0].stats_complete is False # a line was lost, so not a full scan def test_profile_classifies_roles_type_and_verifiability(tmp_path): @@ -436,7 +444,7 @@ def test_profile_tolerates_non_object_jsonl_lines(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) assert result.partitions[0].splits[0].num_examples == 2 # objects counted, stray array dropped - assert result.sampling.exhaustive is True + assert result.partitions[0].stats_complete is True def test_profile_survives_a_hostile_directory(tmp_path): @@ -459,9 +467,9 @@ def test_profile_survives_a_hostile_directory(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise # Nothing here is exhaustive, and the profile says so rather than looking clean. - assert result.sampling.exhaustive is False - assert result.sampling.rows_total is None - assert result.profiler_info["unsupported_files"] == ["leftovers.csv"] + assert result.partitions[0].stats_complete is False + assert result.sampling.rows_present is None + assert [f.path for f in result.unreadable_files] == ["leftovers.csv"] records = {f.path: f for p in result.partitions for s in p.splits for f in s.files} assert records["train-00001-of-00002.parquet"].error is not None # corrupt file, named and explained @@ -512,7 +520,7 @@ def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): records = {f.path: f for p in result.partitions for s in p.splits for f in s.files} assert records["extra.xyz"].num_rows is None # kept, but unreadable - assert result.sampling.exhaustive is False + assert result.partitions[0].stats_complete is False def test_measure_infers_from_rows_when_some_files_declared_no_schema(): @@ -529,3 +537,36 @@ def test_measure_infers_from_rows_when_some_files_declared_no_schema(): features, _, _ = _measure(rows, [declared], exhaustive=True, all_declared=True) assert [f.name for f in features] == ["prompt"] # declared schema trusted when it covers everything + + +def test_stats_completeness_is_per_partition(tmp_path): + # A corrupt shard in one partition says nothing about the measurements in another, but a + # fileset-wide flag downgraded every partition to the worst one. It was never even the value + # that gated quoting a proven enumeration -- that was decided per partition and never stored. + rows = [{"tag": t} for t in ("a", "b", "a")] + _write_parquet(tmp_path / "main" / "train.parquet", rows) + _write_parquet(tmp_path / "socratic" / "train.parquet", rows) + (tmp_path / "socratic" / "broken.parquet").write_bytes(b"not a parquet file") + + partitions = {p.name: p for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions} + + assert partitions["main"].stats_complete is True + assert partitions["socratic"].stats_complete is False + # ...and it is the bit that decides whether a proven enumeration may be quoted. + assert partitions["main"].stats["tag"].categorical.values == ["a", "b"] + assert partitions["socratic"].stats["tag"].categorical.values is None + + +def test_dataset_wide_completeness_is_one_expression(tmp_path): + # SamplingInfo no longer carries `exhaustive`; the contract documents this derivation in its + # place. It has to keep working, or dropping the flag cost consumers something -- and it now + # says *which* half failed, which the single bit could not. + _write_parquet(tmp_path / "train.parquet", [{"a": 1}]) + clean = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + assert all(p.stats_complete for p in clean.partitions) and not clean.unreadable_files + + (tmp_path / "extra.csv").write_text("a,b\n1,2\n") + with_csv = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + assert all(p.stats_complete for p in with_csv.partitions) # the parquet rows are still complete + assert with_csv.unreadable_files # but there is data here that went unprofiled + assert with_csv.sampling.files_read == 1 and with_csv.sampling.files_present == 2 From 51375f82ecfe3c89e274032bcb551829c74930c3 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 5 Aug 2026 11:37:59 -0400 Subject: [PATCH 23/44] fix(datasets): gate quoting a column's values on its role, not its cardinality categorical.values was gated on `exhaustive and distinct_count <= 32`, which tests cardinality while trying to prove enumeration. Those come apart completely on small data: in a three-row dataset every column holds under 32 distinct values, free text included, so an exhaustive scan of a tiny fileset stored an entire column of prompts verbatim -- in a profile the contract advertises as safe to display and export. Checked the authz map before treating this as an escalation: filesets.profile.read and filesets.read are both on Viewer, so nobody sees data they could not already fetch. The exposure is that profiles surface in UIs, logs and exports where the raw rows would not. A ratio (distinct_count / rows_scanned) fixes the statistical error but not the question. 100 clinical notes with 20 distinct diagnoses passes any ratio, genuinely is an enumeration, and still quotes "Stage III breast carcinoma". The profiler cannot read sensitivity out of content. So gate on the signal it does have: the role. label, provenance, meta and rank are controlled vocabularies by construction, at any dataset size; prompt, completion, chosen, rejected, context and messages are free text no matter how few distinct values a small sample shows. An allowlist, not a denylist, so an unroled column fails to silence rather than to exposure. This inverts the order of the measure stage. Roles do not exist when stats are derived -- classification needs the stats first, to tell a binary label column from a class index -- so derive_stats no longer quotes at all and quote_enumerations fills values in afterwards. Filling in rather than redacting is deliberate: skip the pass and nothing is stored, where a skipped redaction would leak. Dropped the ratio guard I had planned as a secondary check. It is redundant with the 32-value cap: a provenance column holding 10k URLs fails the cap already, and a ratio is noisy on exactly the small samples where the cap is doing the work. `exhaustive` leaves this path entirely. It was gating two things -- may I quote this, and is this the complete set -- and only the first is a permission. The second is now expressible: stats_complete says whether the vocabulary is whole, so "observed values (sample)" can be reported instead of withheld. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 19 +++- .../nemo_datasets_plugin/profiler/pipeline.py | 15 +-- .../nemo_datasets_plugin/profiler/stats.py | 74 +++++++++---- plugins/nemo-datasets/tests/test_pipeline.py | 13 +-- plugins/nemo-datasets/tests/test_stats.py | 100 +++++++++++++----- 5 files changed, 155 insertions(+), 66 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 059cd3bc25..a953d686ba 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -214,9 +214,10 @@ def _fields_and_items_are_exclusive(self) -> FeatureSchema: class CategoricalStats(BaseModel): """Cardinality signals for string / int columns. - ``distinct_count`` is always safe to store; the values themselves ARE row data, so they appear - only when proven to be a small enumeration by an exhaustive scan — the same - assert-only-what-was-proven rule applied everywhere the profiler would otherwise leak row content. + ``distinct_count`` is a count, not row content, and is always safe to store. The values + themselves ARE row content, so they appear only for a column whose detected role makes it a + controlled vocabulary — the assert-only-what-was-proven rule applied to the one place the + profiler would otherwise leak the data it is describing. """ distinct_count: int = Field( @@ -227,7 +228,14 @@ class CategoricalStats(BaseModel): ) values: list[str] | None = Field( default=None, - description="The proven enumeration; only when the scan was exhaustive and distinct_count <= 32.", + description=( + "The observed values, present only when this column's `semantic_role` marks it a controlled " + "vocabulary (label | provenance | meta | rank) and it holds at most 32 of them. Cardinality " + "alone cannot be the gate: it inverts on small data, where every column holds few distinct " + "values — free text included — so a three-row dataset had its prompts stored verbatim. A role " + "says what a column *is*, at any size. Read `PartitionProfile.stats_complete` to know whether " + "this is the whole vocabulary or only what the sampled rows showed." + ), ) @@ -237,7 +245,8 @@ class ColumnStats(BaseModel): The kind-specific block is populated by dtype; deep measurements fold into it (e.g. ``MessageStats.content_chars``) so stats stay flat — no path addressing to drift against the schema tree. Never row values — profiles stay safe to display / export without leaking data — - with one gated exception: ``categorical.values``, a proven small enumeration. + with one role-gated exception: ``categorical.values``, and only for a column whose role makes it + a controlled vocabulary rather than free text that happens to repeat. """ null_rate: float = Field(default=0.0, ge=0.0, le=1.0) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 8213883379..4de7da7cc7 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -28,7 +28,7 @@ from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import resolve_splits -from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats +from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats, quote_enumerations from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, DatasetProfile, @@ -184,7 +184,6 @@ def _measure( partition_rows: list[dict], arrow_schemas: list[pa.Schema], *, - exhaustive: bool, all_declared: bool, ) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: """Derive schema, stats and classification, degrading to structure-only if any of it fails. @@ -205,11 +204,15 @@ def _measure( try: declared = _unify_schemas(arrow_schemas) if all_declared else None features = derive_features(partition_rows, declared) - stats = derive_stats(features, partition_rows, exhaustive=exhaustive) + stats = derive_stats(features, partition_rows) # Probes are measured over every column, independent of the roles classify is about to # assign, so a content signal survives a column name the alias table does not know. probes = derive_probes(features, partition_rows) - return features, stats, classify(features, stats, partition_rows, probes=probes) + classification = classify(features, stats, partition_rows, probes=probes) + # Last, because the roles classification assigns are what decide whether a column's values + # may be quoted at all — cardinality only bounds how many. + quote_enumerations(features, stats, partition_rows) + return features, stats, classification except Exception as exc: detail = f"could not measure this partition: {type(exc).__name__}: {exc}" return [], {}, PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]) @@ -310,9 +313,7 @@ def _profile_partition( num_examples=split_examples if split_counts_known else None, ) ) - features, stats, classification = _measure( - partition_rows, arrow_schemas, exhaustive=partition_scanned, all_declared=all_declared - ) + features, stats, classification = _measure(partition_rows, arrow_schemas, all_declared=all_declared) partition = PartitionProfile( name=name, source_dir=source_dir, diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 2d3a7b115b..9a3acaf03f 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -6,8 +6,8 @@ Given a partition's features and its sampled rows, measure each top-level column according to its dtype: length quantiles and corruption signals for text, min/max/mean for numbers, chat-shape signals for messages, and cardinality for both. The result is sparse — a column with nothing worth -measuring is omitted. Row values themselves are never stored, except a proven small enumeration -under ``categorical.values`` when the read was exhaustive. +measuring is omitted. Row values themselves are never stored here at all; a small controlled +vocabulary is added afterwards by :func:`quote_enumerations`, which gates on the column's role. :func:`derive_probes` additionally reads each column's *content* — answer markers, embedded transcripts — as plain per-column counts. Those are measurements, not interpretations: what they @@ -33,14 +33,23 @@ TextStats, ) -# A proven enumeration is only stored when the read was exhaustive and this small. +# A quotable enumeration holds at most this many distinct values. _MAX_ENUM_VALUES = 32 +# Roles that are controlled vocabularies by construction, and so are safe to quote at any dataset +# size. Everything else -- prompts, completions, chosen/rejected, context, chat -- is free text no +# matter how few distinct values a small sample happens to show, and unroled columns are unknown, +# which is the same thing for this purpose. An allowlist, so an unrecognized column fails to silence +# rather than to exposure. +_QUOTABLE_ROLES = frozenset({"label", "provenance", "meta", "rank"}) -def derive_stats( - features: list[FeatureSchema], rows: list[dict[str, Any]], *, exhaustive: bool -) -> dict[str, ColumnStats]: - """Measure each top-level column. Keys are a subset of the feature names (sparse).""" + +def derive_stats(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> dict[str, ColumnStats]: + """Measure each top-level column. Keys are a subset of the feature names (sparse). + + Never fills in ``categorical.values``: that needs the roles, which classification has not + assigned yet. :func:`quote_enumerations` adds them afterwards. + """ total = len(rows) stats: dict[str, ColumnStats] = {} for feature in features: @@ -48,13 +57,39 @@ def derive_stats( # skipping the rest makes which one wins deterministic instead of "whichever came last". if feature.name in stats: continue - column = _column_stats(feature, [row.get(feature.name) for row in rows], total, exhaustive) + column = _column_stats(feature, [row.get(feature.name) for row in rows], total) if column is not None: stats[feature.name] = column return stats -def _column_stats(feature: FeatureSchema, values: list[Any], total: int, exhaustive: bool) -> ColumnStats | None: +def quote_enumerations( + features: list[FeatureSchema], stats: dict[str, ColumnStats], rows: list[dict[str, Any]] +) -> None: + """Fill in ``categorical.values`` for columns whose role makes them a controlled vocabulary. + + Runs after classification, because it needs the roles it gates on, and mutates ``stats`` in place + the way classification mutates ``features``. Deliberately fills in rather than redacting: skip + this pass and no values are stored, where a redaction pass that got skipped would leak them. + + Cardinality is only the size bound. It cannot be the permission, because it inverts on small + data -- in a three-row dataset every column holds under 32 distinct values, free text included, + so an entire column of prompts was quotable. The role says what a column *is*, at any size. + """ + for feature in features: + if feature.semantic_role not in _QUOTABLE_ROLES: + continue + column = stats.get(feature.name) + if column is None or column.categorical is None or column.categorical.distinct_count > _MAX_ENUM_VALUES: + continue + try: + distinct = {value for row in rows if (value := row.get(feature.name)) is not None} + except TypeError: + continue # unhashable values have no enumeration to quote + column.categorical.values = sorted(str(value) for value in distinct) + + +def _column_stats(feature: FeatureSchema, values: list[Any], total: int) -> ColumnStats | None: present = [value for value in values if value is not None] null_rate = (total - len(present)) / total if total else 0.0 @@ -65,14 +100,14 @@ def _column_stats(feature: FeatureSchema, values: list[Any], total: int, exhaust text = TextStats(chars=_quantiles([len(value) for value in strings])) quality = _text_quality(strings) # distinct_count is always safe to store and is the id-like signal (~= rows_scanned) the - # contract documents; only the values themselves are row data, and _cardinality already gates - # those on an exhaustive read. Withholding the count for high-cardinality strings dropped the - # signal precisely where it carries the most information. - categorical = _cardinality(present, exhaustive) + # contract documents. Only the values themselves are row data, and those are added later, + # by role. Withholding the count for high-cardinality strings dropped the signal precisely + # where it carries the most information. + categorical = _cardinality(present) elif feature.dtype == "bool": # The column that decides unpaired_preference deserves a measured class balance rather than # no stats at all. - categorical = _cardinality(present, exhaustive) + categorical = _cardinality(present) elif _is_numeric(feature.dtype): # Drop non-finite floats (NaN / +-inf): they serialize to JSON null and then fail to # re-validate against NumericStats' required floats, which would make the whole profile @@ -84,7 +119,7 @@ def _column_stats(feature: FeatureSchema, values: list[Any], total: int, exhaust ] if numbers: numeric = NumericStats(min=min(numbers), max=max(numbers), mean=sum(numbers) / len(numbers)) - categorical = _cardinality(present, exhaustive) + categorical = _cardinality(present) elif feature.dtype == "messages": messages = _message_stats([value for value in present if isinstance(value, list)]) @@ -114,15 +149,14 @@ def at(percentile: int) -> int: return Quantiles(p50=at(50), p95=at(95), p99=at(99), max=ordered[-1] if ordered else 0) -def _cardinality(present: list[Any], exhaustive: bool) -> CategoricalStats | None: +def _cardinality(present: list[Any]) -> CategoricalStats | None: + """The count only. The values themselves are row data and are gated on role, not cardinality, + so :func:`quote_enumerations` adds them once classification has assigned one.""" try: distinct = set(present) except TypeError: return None # unhashable values (dicts / lists) have no cardinality signal - values = None - if exhaustive and len(distinct) <= _MAX_ENUM_VALUES: - values = sorted(str(value) for value in distinct) - return CategoricalStats(distinct_count=len(distinct), values=values) + return CategoricalStats(distinct_count=len(distinct)) # --- text quality -------------------------------------------------------------------------------- diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index fd0d859ed8..1493ee9a6e 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -531,11 +531,11 @@ def test_measure_infers_from_rows_when_some_files_declared_no_schema(): declared = pa.schema([pa.field("prompt", pa.string())]) rows = [{"prompt": "a"}, {"prompt": "b", "extra": "only in the schemaless file"}] - features, stats, _ = _measure(rows, [declared], exhaustive=True, all_declared=False) + features, stats, _ = _measure(rows, [declared], all_declared=False) assert [f.name for f in features] == ["prompt", "extra"] # the sole witness survives assert set(stats) <= {f.name for f in features} - features, _, _ = _measure(rows, [declared], exhaustive=True, all_declared=True) + features, _, _ = _measure(rows, [declared], all_declared=True) assert [f.name for f in features] == ["prompt"] # declared schema trusted when it covers everything @@ -543,7 +543,7 @@ def test_stats_completeness_is_per_partition(tmp_path): # A corrupt shard in one partition says nothing about the measurements in another, but a # fileset-wide flag downgraded every partition to the worst one. It was never even the value # that gated quoting a proven enumeration -- that was decided per partition and never stored. - rows = [{"tag": t} for t in ("a", "b", "a")] + rows = [{"label": t} for t in (True, False, True)] _write_parquet(tmp_path / "main" / "train.parquet", rows) _write_parquet(tmp_path / "socratic" / "train.parquet", rows) (tmp_path / "socratic" / "broken.parquet").write_bytes(b"not a parquet file") @@ -552,9 +552,10 @@ def test_stats_completeness_is_per_partition(tmp_path): assert partitions["main"].stats_complete is True assert partitions["socratic"].stats_complete is False - # ...and it is the bit that decides whether a proven enumeration may be quoted. - assert partitions["main"].stats["tag"].categorical.values == ["a", "b"] - assert partitions["socratic"].stats["tag"].categorical.values is None + # Quoting is decided by role, not by completeness, so both keep their label vocabulary -- + # stats_complete is what tells a consumer whether socratic's list is the whole of it. + assert partitions["main"].stats["label"].categorical.values == ["False", "True"] + assert partitions["socratic"].stats["label"].categorical.values == ["False", "True"] def test_dataset_wide_completeness_is_one_expression(tmp_path): diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index b764909f61..a5841cddae 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -3,7 +3,7 @@ """Tests for per-column statistics.""" -from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats +from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats, quote_enumerations from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema @@ -20,7 +20,7 @@ def _rows(name, values): def test_text_stats_length_quantiles_and_quality(): values = ["a", "bb", "ccc", "dddd"] - stats = derive_stats([_feature("t", "string")], _rows("t", values), exhaustive=False)["t"] + stats = derive_stats([_feature("t", "string")], _rows("t", values))["t"] assert stats.text.chars.max == 4 assert stats.text.chars.p50 in {2, 3} # nearest-rank over 4 values assert stats.quality is not None @@ -28,56 +28,60 @@ def test_text_stats_length_quantiles_and_quality(): def test_text_quality_flags_repetition_and_non_ascii(): - stats = derive_stats([_feature("t", "string")], _rows("t", ["aaaaaaaa", "héllo wörld"]), exhaustive=False)["t"] + stats = derive_stats([_feature("t", "string")], _rows("t", ["aaaaaaaa", "héllo wörld"]))["t"] assert stats.quality.repetition_score > 0.0 # the "aaaaaaaa" run assert stats.quality.non_ascii_ratio > 0.0 # accented characters -def test_string_cardinality_counts_always_but_withholds_free_text_values(): - # distinct_count is the id-like signal and is always safe to store; only the values themselves - # are row data, and only a small proven enumeration may be kept. - free_text = derive_stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)]), exhaustive=True) +def test_cardinality_counts_are_always_stored(): + # distinct_count is a count, not row content, so it is always safe -- and it is the id-like + # signal the contract documents. + free_text = derive_stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)])) assert free_text["t"].categorical.distinct_count == 50 # ~= row count -> id-like - assert free_text["t"].categorical.values is None # too many distinct values to be an enumeration - labels = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"]), exhaustive=True) + labels = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"])) assert labels["c"].categorical.distinct_count == 2 - assert labels["c"].categorical.values == ["no", "yes"] # proven enumeration under exhaustive read + + +def test_derive_stats_never_quotes_values(): + # Quoting needs a role, and roles are not assigned when stats are measured. Filling them in + # afterwards rather than redacting means a skipped pass stores nothing instead of leaking. + stats = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no"])) + assert stats["c"].categorical.values is None def test_bool_column_gets_a_measured_class_balance(): - stats = derive_stats([_feature("label", "bool")], _rows("label", [True, False, True]), exhaustive=True) + stats = derive_stats([_feature("label", "bool")], _rows("label", [True, False, True])) assert stats["label"].categorical.distinct_count == 2 - assert stats["label"].categorical.values == ["False", "True"] # --- numeric ------------------------------------------------------------------------------------- def test_numeric_stats_and_cardinality(): - stats = derive_stats([_feature("n", "int64")], _rows("n", [0, 4, 2, 2, 3]), exhaustive=True)["n"] + stats = derive_stats([_feature("n", "int64")], _rows("n", [0, 4, 2, 2, 3]))["n"] assert (stats.numeric.min, stats.numeric.max) == (0.0, 4.0) assert stats.numeric.mean == 2.2 assert stats.categorical.distinct_count == 4 # {0, 2, 3, 4} -def test_numeric_cardinality_values_withheld_when_not_exhaustive(): - stats = derive_stats([_feature("n", "int64")], _rows("n", [1, 2, 3]), exhaustive=False)["n"] +def test_numeric_cardinality_counts_without_quoting(): + stats = derive_stats([_feature("n", "int64")], _rows("n", [1, 2, 3]))["n"] assert stats.categorical.distinct_count == 3 - assert stats.categorical.values is None # a sample cannot prove the enumeration + assert stats.categorical.values is None def test_numeric_stats_ignore_non_finite_values(): # NaN / +-inf poison min/max/mean and serialize to JSON null, which then fails to re-validate # against NumericStats' required floats -- making the whole profile unreadable. Drop them. values = [1.0, float("nan"), 3.0, float("inf"), float("-inf"), 5.0] - stats = derive_stats([_feature("n", "float64")], _rows("n", values), exhaustive=True)["n"] + stats = derive_stats([_feature("n", "float64")], _rows("n", values))["n"] assert (stats.numeric.min, stats.numeric.max, stats.numeric.mean) == (1.0, 5.0, 3.0) ColumnStats.model_validate_json(stats.model_dump_json()) # round-trips: no NaN/inf leaked into JSON def test_numeric_all_non_finite_yields_no_numeric_summary(): - stats = derive_stats([_feature("n", "float64")], _rows("n", [float("nan"), float("inf")]), exhaustive=True) + stats = derive_stats([_feature("n", "float64")], _rows("n", [float("nan"), float("inf")])) assert stats.get("n") is None or stats["n"].numeric is None @@ -89,7 +93,7 @@ def test_message_stats_shape_signals(): {"m": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello there"}]}, {"m": [{"role": "user", "content": "again"}, {"role": "assistant", "content": "yes"}]}, ] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.turns.max == 2 assert stats.messages.roles_seen == ["user", "assistant"] # first-seen order assert stats.messages.ends_with_assistant_rate == 1.0 @@ -99,20 +103,20 @@ def test_message_stats_shape_signals(): def test_message_stats_detects_tool_calls_and_user_ending(): rows = [{"m": [{"role": "user", "content": "run"}, {"role": "assistant", "tool_calls": [{"id": "1"}]}]}] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.has_tool_calls is True assert stats.messages.ends_with_assistant_rate == 1.0 # last turn is the assistant tool call def test_message_ends_with_user_turn_is_prompt_only_signal(): rows = [{"m": [{"role": "user", "content": "solve"}]}] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.ends_with_assistant_rate == 0.0 def test_message_stats_read_sharegpt_from_value(): rows = [{"m": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello there"}]}] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.roles_seen == ["human", "gpt"] # verbatim, not normalized assert stats.messages.content_chars.max == len("hi") + len("hello there") assert stats.messages.ends_with_assistant_rate == 1.0 # "gpt" is the responder turn @@ -122,7 +126,7 @@ def test_assistant_equivalent_roles_count_as_the_training_target(): # Matching only the literal "assistant" made every other convention look prompt-only. for responder in ("assistant", "gpt", "bot", "model", "AI"): rows = [{"m": [{"role": "user", "content": "q"}, {"role": responder, "content": "a"}]}] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.ends_with_assistant_rate == 1.0, responder @@ -130,7 +134,7 @@ def test_non_string_role_does_not_break_measurement(): # roles_seen is typed list[str]; a numeric role used to raise a ValidationError from inside the # one stage the pipeline did not guard, aborting the whole profile. rows = [{"m": [{"role": 1, "content": "hi"}]}] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.roles_seen == ["1"] @@ -138,14 +142,14 @@ def test_declared_but_unset_tool_calls_is_not_tool_use(): # parquet materializes every declared struct field, so `"tool_calls" in message` reported tool # use for any schema that merely declares the field. rows = [{"m": [{"role": "user", "content": "hi", "tool_calls": None}]}] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.has_tool_calls is False def test_message_content_parts_tolerate_non_string_text(): # A VLM-style content part whose "text" key is present but not a string must not crash measurement. rows = [{"m": [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": None}]}]}] - stats = derive_stats([_feature("m", "messages")], rows, exhaustive=False)["m"] + stats = derive_stats([_feature("m", "messages")], rows)["m"] assert stats.messages.content_chars.max == 0 # no measurable text, and no crash @@ -155,11 +159,11 @@ def test_message_content_parts_tolerate_non_string_text(): def test_unmeasured_dtypes_are_omitted(): features = [_feature("s", "struct"), _feature("j", "json")] rows = [{"s": {"a": 1}, "j": object()}] - assert derive_stats(features, rows, exhaustive=False) == {} + assert derive_stats(features, rows) == {} def test_null_rate_is_reported(): - stats = derive_stats([_feature("t", "string")], _rows("t", ["a", None, "c", None]), exhaustive=False)["t"] + stats = derive_stats([_feature("t", "string")], _rows("t", ["a", None, "c", None]))["t"] assert stats.null_rate == 0.5 @@ -209,3 +213,43 @@ def test_probes_detect_embedded_transcripts(): rows = [{"c": "\n\nHuman: hi\n\nAssistant: hello"}, {"c": "plain prose"}] probes = derive_probes([_feature("c", "string")], rows) assert probes["c"].transcript_marker == 1 + + +# --- quoting a controlled vocabulary --------------------------------------------------------------- + + +def _quoted(name, dtype, values, role): + """Run the real two-step: measure, then quote by role, and report what was stored.""" + feature = _feature(name, dtype) + feature.semantic_role = role + rows = _rows(name, values) + stats = derive_stats([feature], rows) + quote_enumerations([feature], stats, rows) + return stats[name].categorical.values + + +def test_quotes_a_controlled_vocabulary_role(): + assert _quoted("label", "bool", [True, False, True], "label") == ["False", "True"] + assert _quoted("source", "string", ["gsm8k", "math", "gsm8k"], "provenance") == ["gsm8k", "math"] + assert _quoted("category", "string", ["code", "math"], "meta") == ["code", "math"] + + +def test_refuses_to_quote_free_text_however_few_distinct_values(): + # The failure the cardinality gate could not see: in a tiny dataset every column holds under the + # cap, so a whole column of prompts was quotable and the profile stored it verbatim. + rows = ["Patient Alice, SSN 123-45-6000", "Patient Bob, SSN 123-45-6001"] + assert _quoted("prompt", "string", rows, "prompt") is None + assert _quoted("completion", "string", rows, "completion") is None + assert _quoted("chosen", "string", rows, "chosen") is None + + +def test_refuses_to_quote_an_unroled_column(): + # An unrecognized column is unknown, which here is the same as free text: an allowlist means it + # fails to silence rather than to exposure. + assert _quoted("mystery", "string", ["a", "b"], None) is None + + +def test_refuses_to_quote_a_vocabulary_larger_than_the_cap(): + # Role grants permission; cardinality still bounds the size, so a provenance column holding a + # URL list is not mistaken for an enumeration. + assert _quoted("source", "string", [f"src-{i}" for i in range(40)], "provenance") is None From c292a8b750bef94c9146fcd22406721410c613d2 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 5 Aug 2026 11:50:04 -0400 Subject: [PATCH 24/44] feat(datasets): report every dataset type the roles satisfy, and accept role hints Two changes to the same stage, both about the classification layer asserting less than it knew. _detect_type was an ordered if-chain returning the first match. prompt + completion + score + label is genuinely both scored_response and unpaired_preference, and which one a consumer saw depended on which rule sat higher in the function. The chain now collects: `candidates` lists every structure the roles satisfy, most specific first, so candidates[0] == dataset_type and the summary is a projection rather than a coin flip. The chain was already computing all the predicates and throwing them away. dataset_type is documented as a summary rather than a decision input. The semantic_role markers are the basis a consumer should match on. Deliberately not a capability list ("supports DPO"): trainer requirements shift and differ per framework, and the contract already says it provides the basis and stops. Collecting rather than returning early needed one guard. prompt_only means a prompt with nothing to predict, so it is gated on there being no completion / chosen / rejected -- otherwise a prompt_completion set would claim it too and assert the opposite of what the data holds. The second change gives the role table an input channel. _ALIAS_ROLES is ~35 hardcoded English names with no way to say "my `q` column is the prompt", and its misses are silent: a naming mismatch used to zero the whole classification layer. `column_roles` lets a caller declare one, exposed as `nemo datasets profile --column-role q=prompt`. A hint says which column, not what the data is, so the dtype gates still apply -- factored out of _role_for into _dtype_allows and now shared by both paths. A rejected hint is reported as user_hint evidence naming the column and dtype, then falls through to detection, so a bad hint costs nothing the table would have found. Accepting hints unconditionally would let one typo produce a nonsense classification and make the profile a place to store mistakes. FeatureSchema gains semantic_role_source (detected | declared). The distinction is per-column and actionable -- a UI renders a declared role as confirmed and a detected one as a suggestion -- so it is a field rather than something to recover by parsing evidence prose. Reading column_roles from fileset metadata is the platform half and lands with the Files integration; this is the profiler half, with the CLI as its caller. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 33 ++++- .../tests/files/test_dataset_profile.py | 32 +++-- .../src/nemo_datasets_plugin/cli.py | 18 ++- .../nemo_datasets_plugin/profiler/classify.py | 124 +++++++++++++----- .../nemo_datasets_plugin/profiler/pipeline.py | 21 ++- plugins/nemo-datasets/tests/test_classify.py | 84 ++++++++++++ plugins/nemo-datasets/tests/test_cli.py | 20 +++ plugins/nemo-datasets/tests/test_pipeline.py | 4 +- 8 files changed, 278 insertions(+), 58 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index a953d686ba..4b19955006 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -50,7 +50,8 @@ class Evidence(BaseModel): kind: str = Field( description=( "column_name | column_dtype | content_probe | split_name | file_name | card_metadata | " - "error — the last for when a detector could not run at all, so an absent finding is " + "user_hint | error — `user_hint` for a caller-supplied column role the data could not " + "support, and `error` for when a detector could not run at all, so an absent finding is " "distinguishable from a finding of absence." ), ) @@ -83,7 +84,25 @@ class PartitionClassification(BaseModel): """ modality: str = Field(default="text", description="text | image_text | audio_text | ...") - dataset_type: str = Field(description="Dataset-type vocabulary (prompt_completion, preference_pair, ...).") + dataset_type: str = Field( + description=( + "Dataset-type vocabulary (prompt_completion, preference_pair, ...). A SUMMARY, not the " + "basis for a decision — it is the most specific single structure the roles satisfy, and a " + "dataset routinely satisfies several. The `semantic_role` markers are what a consumer " + "should match on; `candidates` lists everything this one is a projection of." + ), + ) + candidates: list[str] = Field( + default_factory=list, + description=( + "Every dataset type the assigned roles satisfy, most specific first, so " + "`candidates[0] == dataset_type`. prompt + completion + score + label is genuinely both " + "scored_response and unpaired_preference; reporting only the first made rule order an " + "invisible tie-break and hid that the data supports more than one use. Deliberately not a " + 'capability list ("supports DPO") — trainer requirements shift and differ per framework, ' + "so that mapping belongs in the consumer, computed from the roles." + ), + ) format: str | None = Field(default=None, description="standard | conversational | mixed") prompt_form: str | None = Field(default=None, description="explicit | implicit | n/a") verifiability: Verifiability | None = Field( @@ -187,6 +206,16 @@ class FeatureSchema(BaseModel): "message struct's `role` key." ), ) + semantic_role_source: str | None = Field( + default=None, + description=( + "Where `semantic_role` came from: detected | declared. A declared role was supplied by the " + "caller and only accepted because the dtype could carry it; a detected one was inferred from " + "the column name. Kept as a field rather than left to evidence prose because the distinction " + "is per-column and actionable — a UI renders a declared role as confirmed and a detected one " + "as a suggestion to correct." + ), + ) fixed_length: int | None = Field( default=None, description=( diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index ebcd4ebc10..8178d036ea 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -50,9 +50,9 @@ files: [{path: test-00000-of-00001.parquet, size_bytes: 411552, checksum: sha256:02af..., num_rows: 200, file_format: parquet, read_strategy: head}]} features: - - {name: prompt, dtype: messages, semantic_role: prompt, + - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} - - {name: completion, dtype: messages, semantic_role: completion, + - {name: completion, dtype: messages, semantic_role: completion, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} stats: prompt: {messages: {turns: {p50: 1, p95: 1, p99: 1, max: 1}, content_chars: {p50: 180, p95: 620, p99: 1100, max: 4800}, @@ -62,6 +62,7 @@ classification: modality: text dataset_type: prompt_completion + candidates: [prompt_completion] format: conversational prompt_form: explicit verifiability: @@ -93,11 +94,11 @@ files: [{path: test-00000-of-00001.parquet, size_bytes: 1198422, checksum: sha256:5c1d..., num_rows: 2354, file_format: parquet, read_strategy: head}]} features: - - {name: prompt, dtype: messages, semantic_role: prompt, + - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} - - {name: chosen, dtype: messages, semantic_role: chosen, + - {name: chosen, dtype: messages, semantic_role: chosen, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} - - {name: rejected, dtype: messages, semantic_role: rejected, + - {name: rejected, dtype: messages, semantic_role: rejected, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} stats: prompt: {messages: {turns: {p50: 3, p95: 8, p99: 9, max: 9}, content_chars: {p50: 640, p95: 3200, p99: 5400, max: 9800}, @@ -109,6 +110,7 @@ classification: modality: text dataset_type: preference_pair + candidates: [preference_pair] format: conversational prompt_form: explicit evidence: @@ -136,13 +138,13 @@ files: [{path: validation-00000-of-00001.parquet, size_bytes: 2311008, checksum: sha256:8bd2..., num_rows: 1038, file_format: parquet, read_strategy: head}]} features: - - {name: prompt, dtype: string, semantic_role: prompt} - - {name: response, dtype: string, semantic_role: completion} - - {name: helpfulness, dtype: int64, semantic_role: score} - - {name: correctness, dtype: int64, semantic_role: score} - - {name: coherence, dtype: int64, semantic_role: score} - - {name: complexity, dtype: int64, semantic_role: score} - - {name: verbosity, dtype: int64, semantic_role: score} + - {name: prompt, dtype: string, semantic_role: prompt, semantic_role_source: detected} + - {name: response, dtype: string, semantic_role: completion, semantic_role_source: detected} + - {name: helpfulness, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: correctness, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: coherence, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: complexity, dtype: int64, semantic_role: score, semantic_role_source: detected} + - {name: verbosity, dtype: int64, semantic_role: score, semantic_role_source: detected} stats: prompt: {text: {chars: {p50: 320, p95: 2200, p99: 5600, max: 12000}}, quality: {whitespace_ratio: 0.16, non_ascii_ratio: 0.004, repetition_score: 0.02}} @@ -156,6 +158,7 @@ classification: modality: text dataset_type: scored_response + candidates: [scored_response, prompt_completion] format: standard prompt_form: explicit evidence: @@ -216,6 +219,7 @@ def _build_profile() -> DatasetProfile: }, classification=PartitionClassification( dataset_type="prompt_completion", + candidates=["prompt_completion"], format="standard", prompt_form="explicit", verifiability=Verifiability( @@ -275,6 +279,10 @@ def test_helpsteer2_flat_schema_and_no_verifiability(): profile = DatasetProfile.model_validate(yaml.safe_load(HELPSTEER2)) part = profile.partitions[0] assert part.classification.dataset_type == "scored_response" + # A scored prompt/completion set is also a plain prompt_completion set. `dataset_type` is the + # most specific reading; `candidates` is what the same columns otherwise support. + assert part.classification.candidates == ["scored_response", "prompt_completion"] + assert part.classification.candidates[0] == part.classification.dataset_type assert part.classification.format == "standard" # Absence of a verifiability object *is* the "not verifiable" claim. assert part.classification.verifiability is None diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py index 7adf17aa49..577f515c31 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py @@ -36,6 +36,14 @@ def profile( "but scales memory with the dataset rather than the file count.", min=0, ), + column_role: list[str] = typer.Option( + None, + "--column-role", + help="Assert a column's role as NAME=ROLE (repeatable), e.g. --column-role q=prompt. " + "Takes precedence over name detection, but the dtype must still support the role; a " + "rejected hint is reported in the profile's evidence.", + metavar="NAME=ROLE", + ), ) -> None: """Profile a local dataset directory and print its DatasetProfile.""" # Imported here, not at module scope: the platform calls get_cli() for every plugin at @@ -46,15 +54,21 @@ def profile( if output not in {"json", "yaml"}: raise typer.BadParameter("output must be 'json' or 'yaml'") + column_roles: dict[str, str] = {} + for pair in column_role or []: + name, separator, role = pair.partition("=") + if not separator or not name or not role: + raise typer.BadParameter(f"--column-role expects NAME=ROLE, got {pair!r}") + column_roles[name] = role try: source = LocalFileSource(path) except NotADirectoryError as exc: raise typer.BadParameter(str(exc)) from exc if rows_per_file is None: - result = run_profile(source) + result = run_profile(source, column_roles=column_roles) else: - result = run_profile(source, row_cap=rows_per_file or None) + result = run_profile(source, row_cap=rows_per_file or None, column_roles=column_roles) if output == "yaml": import yaml diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index 207a8400a3..9426531c20 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -109,37 +109,71 @@ def _is_label_column(feature: FeatureSchema, stats: dict[str, ColumnStats]) -> b return _is_numeric(feature.dtype) and _is_binary(stats.get(feature.name)) +def _dtype_allows(feature: FeatureSchema, role: str, stats: dict[str, ColumnStats]) -> bool: + """Whether this column's dtype can carry ``role`` at all. + + Applied to detected *and* declared roles alike. A hint says which column, not what the data is: + without this, one typo (``{"score": "prompt"}`` on an int column) would silently produce a + nonsense classification and the profile would become a place to store mistakes. + """ + dtype = feature.dtype + if role == "score" or role == "rank": + return _is_numeric(dtype) + if role == "label": + return _is_label_column(feature, stats) + if role == "messages": + return dtype == "messages" + if role in {"prompt", "completion", "chosen", "rejected", "context", "system"}: + return dtype in _TEXT_DTYPES + if role == "ground_truth": + return dtype in _GROUND_TRUTH_DTYPES + if role in {"stepwise_completions", "stepwise_labels"}: + return dtype == "list" + return True # id / provenance / meta / tools / image carry no dtype constraint + + def _role_for(feature: FeatureSchema, stats: dict[str, ColumnStats]) -> str | None: + """The role this column's *name* implies, if the dtype does not contradict it.""" name = feature.name.lower() - dtype = feature.dtype - if name in _SCORE_ALIASES and _is_numeric(dtype): + if name in _SCORE_ALIASES and _is_numeric(feature.dtype): return "score" - if name == "label": - return "label" if _is_label_column(feature, stats) else None - role = _ALIAS_ROLES.get(name) if role is None: return None - # dtype gates: reject an alias whose dtype contradicts the role. - if role == "messages" and dtype != "messages": - return None - if role in {"prompt", "completion", "chosen", "rejected", "context", "system"}: - if dtype not in _TEXT_DTYPES: - return None - if role == "ground_truth" and dtype not in _GROUND_TRUTH_DTYPES: - return None - if role == "rank" and not _is_numeric(dtype): - return None - if role in {"stepwise_completions", "stepwise_labels"} and dtype != "list": - return None - return role + return role if _dtype_allows(feature, role, stats) else None -def _assign_roles(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> None: +def _assign_roles( + features: list[FeatureSchema], stats: dict[str, ColumnStats], column_roles: dict[str, str] +) -> list[Evidence]: + """Stack roles onto ``features`` in place; return evidence for any hint the data could not support. + + A declared role wins over the name-alias table — the caller knows their schema and the table is + ~35 English names — but it still has to pass the dtype gate, and a rejected hint is reported + rather than dropped. Silence is what made the alias table's misses so expensive in the first place. + """ + rejected: list[Evidence] = [] for feature in features: + declared = column_roles.get(feature.name) + if declared is not None: + if _dtype_allows(feature, declared, stats): + feature.semantic_role = declared + feature.semantic_role_source = "declared" + continue + rejected.append( + Evidence( + kind="user_hint", + detail=( + f"hint '{feature.name} -> {declared}' rejected: a {feature.dtype} column " + f"cannot carry that role; falling back to detection" + ), + ) + ) role = _role_for(feature, stats) if role is not None: feature.semantic_role = role + feature.semantic_role_source = "detected" + return rejected def _detect_modality(features: list[FeatureSchema]) -> str: @@ -178,37 +212,49 @@ def _messages_stats(features: list[FeatureSchema], stats: dict[str, ColumnStats] return None -def _detect_type(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> str: +def _detect_types(features: list[FeatureSchema], stats: dict[str, ColumnStats]) -> list[str]: + """Every dataset type the assigned roles satisfy, most specific first. + + The chain is ordered by specificity, so the head is the best single answer and the tail is + structures the same columns *also* satisfy. Returning only the head made rule order an invisible + tie-break: prompt + completion + score + label is genuinely both scored_response and + unpaired_preference, and which one a consumer saw depended on line numbers. + """ roles = {feature.semantic_role for feature in features if feature.semantic_role} + targets = roles & {"completion", "chosen", "rejected", "stepwise_completions"} + candidates: list[str] = [] def has(*required: str) -> bool: return all(role in roles for role in required) if has("prompt", "stepwise_completions", "stepwise_labels"): - return "stepwise_supervision" + candidates.append("stepwise_supervision") if has("chosen", "rejected"): - return "preference_pair" + candidates.append("preference_pair") if has("prompt", "completion", "score"): - return "scored_response" + candidates.append("scored_response") if has("prompt", "completion", "label"): - return "unpaired_preference" + candidates.append("unpaired_preference") # `rank` is only a dataset type alongside something to rank. On its own it short-circuited every # more specific structure above, so a stray numeric column named `rank` — or a ranked variant of # a preference set — was enough to mislabel the dataset. - if has("rank") and roles & {"completion", "chosen", "rejected", "stepwise_completions"}: - return "ranked_responses" + if has("rank") and targets: + candidates.append("ranked_responses") if has("prompt", "completion"): - return "prompt_completion" + candidates.append("prompt_completion") if "messages" in roles: message_stats = _messages_stats(features, stats) if message_stats is not None and message_stats.ends_with_assistant_rate < 0.5: - return "prompt_only" # a chat that ends on a user turn has no training target - return "messages" - if "prompt" in roles: - return "prompt_only" + candidates.append("prompt_only") # a chat that ends on a user turn has no training target + else: + candidates.append("messages") + # A prompt with nothing to predict. Guarded on `targets` because with candidates collected rather + # than returned early, a prompt+completion set would otherwise claim prompt_only as well. + if "prompt" in roles and not targets and "prompt_only" not in candidates: + candidates.append("prompt_only") if len(features) == 1 and features[0].dtype == "string" and features[0].semantic_role is None: - return "text" - return "unknown" + candidates.append("text") + return candidates or ["unknown"] # --- interpreting the content probes -------------------------------------------------------------- @@ -318,6 +364,7 @@ def classify( rows: list[dict] | None = None, *, probes: dict[str, ColumnProbes] | None = None, + column_roles: dict[str, str] | None = None, ) -> PartitionClassification: """Assign roles onto ``features`` in place and return the partition's classification. @@ -325,16 +372,20 @@ def classify( a pure function of ``(features, rows)``, so a caller that has not already computed them can pass ``rows`` alone and get them derived here; the pipeline passes them in to avoid the second pass. Role/axis/type inference needs neither — only the schema and stats. + + ``column_roles`` maps a column name to a role the caller is asserting, taking precedence over + the name-alias table but still subject to the dtype gates. It exists because that table is ~35 + English names with no way to say "my `q` column is the prompt", and its misses are silent. """ rows = rows or [] probes = derive_probes(features, rows) if probes is None else probes - _assign_roles(features, stats) + evidence = _assign_roles(features, stats, column_roles or {}) roles = {feature.semantic_role for feature in features if feature.semantic_role} - dataset_type = _detect_type(features, stats) + candidates = _detect_types(features, stats) + dataset_type = candidates[0] fmt = _detect_format(features) prompt_form = _detect_prompt_form(roles) if dataset_type != "unknown" else None - evidence: list[Evidence] = [] role_columns = [f"{feature.name} -> {feature.semantic_role}" for feature in features if feature.semantic_role] if role_columns: evidence.append(Evidence(kind="column_name", detail=f"columns matched roles: {', '.join(role_columns)}")) @@ -348,6 +399,7 @@ def classify( return PartitionClassification( modality=_detect_modality(features), dataset_type=dataset_type, + candidates=candidates, format=fmt, prompt_form=prompt_form, verifiability=_detect_verifiability(features, probes), diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 4de7da7cc7..48ae736db1 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -67,6 +67,7 @@ def profile( *, created_at: datetime | None = None, row_cap: int | None = DEFAULT_ROW_CAP, + column_roles: dict[str, str] | None = None, ) -> DatasetProfile: """Profile the dataset behind ``source`` into a ``DatasetProfile``. @@ -74,6 +75,10 @@ def profile( exact but scales memory with the dataset rather than the file count. Files smaller than the cap are still read to the end, so a capped profile of a small dataset stays exhaustive. + ``column_roles`` maps a column name to a role the caller is asserting, for datasets whose column + names the role table does not recognize. Hints take precedence over name detection but still have + to pass the dtype gates, and a rejected one is reported as evidence rather than dropped. + ``created_at`` is injectable so a profile can be made reproducible byte-for-byte in tests; it defaults to the current UTC time. """ @@ -105,7 +110,7 @@ def profile( groups = group_partitions(data_entries) for source_dir, partition_entries in groups: name = _partition_label(source_dir, len(groups)) - outcome = _profile_partition(source, name, source_dir, partition_entries, row_cap) + outcome = _profile_partition(source, name, source_dir, partition_entries, row_cap, column_roles or {}) partitions.append(outcome.partition) rows_scanned += outcome.rows_scanned files_read += outcome.files_read @@ -185,6 +190,7 @@ def _measure( arrow_schemas: list[pa.Schema], *, all_declared: bool, + column_roles: dict[str, str], ) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: """Derive schema, stats and classification, degrading to structure-only if any of it fails. @@ -208,7 +214,7 @@ def _measure( # Probes are measured over every column, independent of the roles classify is about to # assign, so a content signal survives a column name the alias table does not know. probes = derive_probes(features, partition_rows) - classification = classify(features, stats, partition_rows, probes=probes) + classification = classify(features, stats, partition_rows, probes=probes, column_roles=column_roles) # Last, because the roles classification assigns are what decide whether a column's values # may be quoted at all — cardinality only bounds how many. quote_enumerations(features, stats, partition_rows) @@ -229,7 +235,12 @@ class _PartitionOutcome: def _profile_partition( - source: FileSource, name: str, source_dir: str | None, entries: list[FileEntry], row_cap: int | None + source: FileSource, + name: str, + source_dir: str | None, + entries: list[FileEntry], + row_cap: int | None, + column_roles: dict[str, str], ) -> _PartitionOutcome: """Profile one partition — the files of one source directory, whatever formats they are in. @@ -313,7 +324,9 @@ def _profile_partition( num_examples=split_examples if split_counts_known else None, ) ) - features, stats, classification = _measure(partition_rows, arrow_schemas, all_declared=all_declared) + features, stats, classification = _measure( + partition_rows, arrow_schemas, all_declared=all_declared, column_roles=column_roles + ) partition = PartitionProfile( name=name, source_dir=source_dir, diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index 9c84a2174e..05858e7438 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -18,6 +18,11 @@ def _f(name, dtype): return FeatureSchema(name=name, dtype=dtype) +def _binary_column(): + """A column observed to hold two distinct values -- what makes an int/bool a real label.""" + return ColumnStats(categorical=CategoricalStats(distinct_count=2)) + + def _messages_column(ends_with_assistant_rate): q = Quantiles(p50=1, p95=1, p99=1, max=1) return ColumnStats( @@ -295,3 +300,82 @@ def test_precomputed_probes_and_derived_probes_agree(): passed_in = classify(features, {}, rows, probes=derive_probes(features, rows)) assert derived.verifiability.coverage == passed_in.verifiability.coverage assert derived.verifiability.evidence[0].detail == passed_in.verifiability.evidence[0].detail + + +# --- candidates ------------------------------------------------------------------------------------ + + +def test_candidates_list_every_structure_the_roles_satisfy(): + # prompt + completion + score + label is genuinely both scored_response and unpaired_preference. + # Reporting only the first made rule order an invisible tie-break. + features = [_f("prompt", "string"), _f("completion", "string"), _f("score", "float64"), _f("label", "bool")] + result = classify(features, {"label": _binary_column()}) + + assert result.candidates == ["scored_response", "unpaired_preference", "prompt_completion"] + assert result.dataset_type == result.candidates[0] # the summary is the head, never more + + +def test_candidates_collapse_to_one_when_the_structure_is_unambiguous(): + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert result.candidates == ["prompt_completion"] + + +def test_unknown_is_still_reported_as_a_candidate(): + result = classify([_f("foo", "int64"), _f("bar", "int64")], {}) + assert result.dataset_type == "unknown" + assert result.candidates == ["unknown"] + + +def test_prompt_only_is_not_claimed_alongside_a_training_target(): + # Collecting candidates rather than returning early risks a prompt+completion set also claiming + # prompt_only, which asserts the opposite of what the data holds. + result = classify([_f("prompt", "string"), _f("completion", "string")], {}) + assert "prompt_only" not in result.candidates + + +# --- declared roles (hints) ------------------------------------------------------------------------ + + +def test_a_hint_names_a_column_the_alias_table_does_not_know(): + features = [_f("q", "string"), _f("a", "string")] + result = classify(features, {}, column_roles={"q": "prompt", "a": "completion"}) + + assert [(f.semantic_role, f.semantic_role_source) for f in features] == [ + ("prompt", "declared"), + ("completion", "declared"), + ] + assert result.dataset_type == "prompt_completion" + + +def test_a_hint_takes_precedence_over_the_name_alias(): + # The caller knows their schema; the table is ~35 English names. + features = [_f("prompt", "string")] + classify(features, {}, column_roles={"prompt": "context"}) + assert features[0].semantic_role == "context" + assert features[0].semantic_role_source == "declared" + + +def test_a_hint_the_dtype_cannot_support_is_rejected_loudly(): + # A hint says which column, not what the data is. Accepting it unconditionally would let one + # typo produce a nonsense classification, and silence is what made the table's misses costly. + features = [_f("n", "int64")] + result = classify(features, {}, column_roles={"n": "prompt"}) + + assert features[0].semantic_role is None + rejections = [e for e in result.evidence if e.kind == "user_hint"] + assert len(rejections) == 1 + assert "n -> prompt" in rejections[0].detail and "int64" in rejections[0].detail + + +def test_a_rejected_hint_falls_back_to_detection(): + # `answer` is a known alias; a bad hint on it must not cost the role the table would have found. + features = [_f("answer", "string")] + classify(features, {}, column_roles={"answer": "messages"}) # messages needs the messages dtype + assert features[0].semantic_role == "completion" + assert features[0].semantic_role_source == "detected" + + +def test_detected_roles_are_marked_as_detected(): + features = [_f("prompt", "string")] + classify(features, {}) + assert features[0].semantic_role_source == "detected" diff --git a/plugins/nemo-datasets/tests/test_cli.py b/plugins/nemo-datasets/tests/test_cli.py index deb335486a..64d665cb6c 100644 --- a/plugins/nemo-datasets/tests/test_cli.py +++ b/plugins/nemo-datasets/tests/test_cli.py @@ -57,3 +57,23 @@ def test_profile_command_rejects_non_directory(tmp_path): target.write_text("x") result = runner.invoke(_mounted(), ["datasets", "profile", str(target)]) assert result.exit_code != 0 + + +def test_profile_command_accepts_column_role_hints(tmp_path): + # The hint mechanism needs a caller on this branch; reading them from fileset metadata is the + # platform half and lands with the Files integration. + pq.write_table(pa.Table.from_pylist([{"q": "why?", "a": "because"}]), tmp_path / "train.parquet") + result = runner.invoke( + _mounted(), ["datasets", "profile", str(tmp_path), "--column-role", "q=prompt", "--column-role", "a=completion"] + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.stdout) + partition = payload["partitions"][0] + assert partition["classification"]["dataset_type"] == "prompt_completion" + assert [f["semantic_role_source"] for f in partition["features"]] == ["declared", "declared"] + + +def test_profile_command_rejects_a_malformed_column_role(tmp_path): + result = runner.invoke(_mounted(), ["datasets", "profile", str(tmp_path), "--column-role", "no-equals-sign"]) + assert result.exit_code != 0 diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 1493ee9a6e..c9ff6f4acf 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -531,11 +531,11 @@ def test_measure_infers_from_rows_when_some_files_declared_no_schema(): declared = pa.schema([pa.field("prompt", pa.string())]) rows = [{"prompt": "a"}, {"prompt": "b", "extra": "only in the schemaless file"}] - features, stats, _ = _measure(rows, [declared], all_declared=False) + features, stats, _ = _measure(rows, [declared], all_declared=False, column_roles={}) assert [f.name for f in features] == ["prompt", "extra"] # the sole witness survives assert set(stats) <= {f.name for f in features} - features, _, _ = _measure(rows, [declared], all_declared=True) + features, _, _ = _measure(rows, [declared], all_declared=True, column_roles={}) assert [f.name for f in features] == ["prompt"] # declared schema trusted when it covers everything From 9b63e4efad49714bd14dda8fc5a3a5c34eab9cec Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 5 Aug 2026 11:58:34 -0400 Subject: [PATCH 25/44] perf(datasets): budget rows per partition instead of capping them per file The row cap was on the wrong axis. It bounded each file while the cost is per partition, so peak memory scaled with shard count: measured at 1352 bytes/row, a dataset resharded from 100 files to 10,000 took peak heap from 135 MB to 13.5 GB while describing exactly the same rows. Lowering the cap to survive that would have crippled the 10-shard case, where it is the only thing producing enough rows to measure. One knob could not serve both ends. `row_budget` bounds a partition's total and is divided across its files, so the same data costs the same to profile however it is sharded. Measured on the fixture that produced the numbers above: 40 shards and 200 shards both peak at 14 MB, where before 40 shards alone cost 54 MB and climbing. Projected at 10,000 shards, 0.14 GB against 13.5 GB. MIN_ROWS_PER_FILE keeps a floor under the division, which makes the budget a target rather than a ceiling: at 10,000 shards the arithmetic share is one row, too thin to witness a column that only that shard holds. Overshooting is the right trade, because the alternative is sampling a subset of *files*, which hides exactly those columns -- the tier of this problem still unsolved, and the one that will matter first, since 10,000 shards is already 10,000 round trips just to open them. SamplingInfo.per_file_row_cap was the caller's request and is now derived, so it splits: `row_budget` carries the request at the envelope and each FileRecord.row_cap carries the share it actually got. Same move as the previous commits -- the value lands where it is decided. Not doing the streaming-accumulator tier. Its purpose was to avoid retaining rows, and with the budget capping a partition at ~10k rows there are few to retain; what it would buy now is accurate stats over an *uncapped* read, which is an accuracy feature rather than the scaling fix it was scoped as. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 21 +++++++- .../tests/files/test_dataset_profile.py | 21 ++++---- .../src/nemo_datasets_plugin/cli.py | 15 +++--- .../nemo_datasets_plugin/profiler/pipeline.py | 52 +++++++++++++----- plugins/nemo-datasets/tests/test_pipeline.py | 53 ++++++++++++++++--- 5 files changed, 123 insertions(+), 39 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 4b19955006..a4c95463c4 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -323,6 +323,15 @@ class FileRecord(BaseModel): "in the same partition can only be read from the top." ), ) + row_cap: int | None = Field( + default=None, + description=( + "Rows this file's read was bounded to, None when unbounded. Derived from " + "`SamplingInfo.row_budget` divided across the partition's files, so it is per file rather " + "than a global setting: the same budget yields 1000 rows each across ten shards and ten " + "each across a thousand, which is what keeps peak memory flat as a dataset is resharded." + ), + ) num_rows: int | None = Field( default=None, description="Exact only (parquet footer / exhaustive scan), else None.", @@ -507,7 +516,17 @@ class SamplingInfo(BaseModel): "these two to match until scale forces file-level sampling." ), ) - per_file_row_cap: int | None = Field(default=None, description="Cap that bounded per-file reads, if any.") + row_budget: int | None = Field( + default=None, + description=( + "Rows the caller allowed per partition, None for an unbounded read. A budget rather than a " + "per-file cap because the cost is per partition: a per-file cap made peak memory scale with " + "shard count, so the same dataset resharded from 100 files to 10,000 went from megabytes to " + "gigabytes without holding any more data. The per-file cap this produced is on each " + "`FileRecord.row_cap`. Not a hard ceiling: every file is still read at least a few rows, " + "since a file sampled too thinly cannot contribute the columns it alone witnesses." + ), + ) seed: int | None = Field(default=None, description="RNG seed used for row selection, for reproducibility.") diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index 8178d036ea..8915e435d1 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -36,7 +36,7 @@ created_at: 2026-07-08T22:05:12Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 2112, rows_present: 3201061, - files_read: 33, files_present: 33, per_file_row_cap: 64} + files_read: 33, files_present: 33, row_budget: 4096} partitions: - name: default source_dir: null @@ -45,10 +45,10 @@ splits: - {name: train, canonical: train, num_examples: 3200861, files: [{path: train-00000-of-00032.parquet, size_bytes: 193777041, - checksum: sha256:9c1e..., num_rows: 100027, file_format: parquet, read_strategy: head}]} + checksum: sha256:9c1e..., num_rows: 100027, file_format: parquet, read_strategy: head, row_cap: 512}]} - {name: test, canonical: test, num_examples: 200, files: [{path: test-00000-of-00001.parquet, size_bytes: 411552, - checksum: sha256:02af..., num_rows: 200, file_format: parquet, read_strategy: head}]} + checksum: sha256:02af..., num_rows: 200, file_format: parquet, read_strategy: head, row_cap: 512}]} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -80,7 +80,7 @@ created_at: 2026-07-08T22:41:37Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 1024, rows_present: 46189, - files_read: 2, files_present: 2, per_file_row_cap: 512} + files_read: 2, files_present: 2, row_budget: 1024} partitions: - name: default source_dir: null @@ -89,10 +89,10 @@ splits: - {name: train, canonical: train, num_examples: 43835, files: [{path: train-00000-of-00001.parquet, size_bytes: 22105331, - checksum: sha256:77b0..., num_rows: 43835, file_format: parquet, read_strategy: head}]} + checksum: sha256:77b0..., num_rows: 43835, file_format: parquet, read_strategy: head, row_cap: 512}]} - {name: test, canonical: test, num_examples: 2354, files: [{path: test-00000-of-00001.parquet, size_bytes: 1198422, - checksum: sha256:5c1d..., num_rows: 2354, file_format: parquet, read_strategy: head}]} + checksum: sha256:5c1d..., num_rows: 2354, file_format: parquet, read_strategy: head, row_cap: 512}]} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -124,7 +124,7 @@ created_at: 2026-07-09T10:12:45Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 1024, rows_present: 21362, - files_read: 2, files_present: 2, per_file_row_cap: 512} + files_read: 2, files_present: 2, row_budget: 1024} partitions: - name: default source_dir: null @@ -133,10 +133,10 @@ splits: - {name: train, canonical: train, num_examples: 20324, files: [{path: train-00000-of-00001.parquet, size_bytes: 44201991, - checksum: sha256:e410..., num_rows: 20324, file_format: parquet, read_strategy: head}]} + checksum: sha256:e410..., num_rows: 20324, file_format: parquet, read_strategy: head, row_cap: 512}]} - {name: validation, canonical: validation, num_examples: 1038, files: [{path: validation-00000-of-00001.parquet, size_bytes: 2311008, - checksum: sha256:8bd2..., num_rows: 1038, file_format: parquet, read_strategy: head}]} + checksum: sha256:8bd2..., num_rows: 1038, file_format: parquet, read_strategy: head, row_cap: 512}]} features: - {name: prompt, dtype: string, semantic_role: prompt, semantic_role_source: detected} - {name: response, dtype: string, semantic_role: completion, semantic_role_source: detected} @@ -184,7 +184,7 @@ def _build_profile() -> DatasetProfile: rows_present=2048, files_read=2, files_present=2, - per_file_row_cap=512, + row_budget=1024, seed=7, ), partitions=[ @@ -204,6 +204,7 @@ def _build_profile() -> DatasetProfile: checksum="sha256:ab", file_format="parquet", read_strategy="head", + row_cap=512, num_rows=2048, ) ], diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py index 577f515c31..923cded46b 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py @@ -29,11 +29,12 @@ def _root() -> None: def profile( path: str = typer.Argument(..., help="Path to a local directory of dataset files."), output: str = typer.Option("json", "--output", "-o", help="Output format: json | yaml."), - rows_per_file: int = typer.Option( + row_budget: int = typer.Option( None, - "--rows-per-file", - help="Rows to read from each file (default 1000); 0 reads every row, which is exact " - "but scales memory with the dataset rather than the file count.", + "--row-budget", + help="Rows to read per partition, divided across its files (default 10000); 0 reads " + "every row, which is exact but scales memory with the dataset. A budget rather than a " + "per-file cap so peak memory does not grow when a dataset is resharded.", min=0, ), column_role: list[str] = typer.Option( @@ -47,7 +48,7 @@ def profile( ) -> None: """Profile a local dataset directory and print its DatasetProfile.""" # Imported here, not at module scope: the platform calls get_cli() for every plugin at - # startup, and the profiler pulls in pyarrow. The row-cap default lives in the pipeline + # startup, and the profiler pulls in pyarrow. The budget default lives in the pipeline # rather than being restated here, so an unspecified flag simply omits the argument. from nemo_datasets_plugin.profiler.file_source import LocalFileSource from nemo_datasets_plugin.profiler.pipeline import profile as run_profile @@ -65,10 +66,10 @@ def profile( except NotADirectoryError as exc: raise typer.BadParameter(str(exc)) from exc - if rows_per_file is None: + if row_budget is None: result = run_profile(source, column_roles=column_roles) else: - result = run_profile(source, row_cap=rows_per_file or None, column_roles=column_roles) + result = run_profile(source, row_budget=row_budget or None, column_roles=column_roles) if output == "yaml": import yaml diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 48ae736db1..668dad7d3c 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -44,13 +44,19 @@ PROFILER_NAME = "nemo-dataset-profiler" PROFILER_VERSION = "0.1.0" -# Rows read per file by default. Every file is still opened — head-sampling a *subset of files* would -# hide columns that appear only in later shards — but each is capped, so peak memory scales with the -# file count rather than the dataset size. Uncapped, a partition materializes every row of every file -# as Python dicts at roughly 6x the on-disk parquet size, which puts a 10 GB dataset far past any -# reasonable machine. A thousand rows per file is ample for the statistics computed here (length -# quantiles, rates, cardinality); pass ``row_cap=None`` for a genuinely exhaustive scan. -DEFAULT_ROW_CAP = 1000 +# Rows a partition may read, in total, by default. Every file is still opened — head-sampling a +# *subset of files* would hide columns that appear only in later shards — but the budget is divided +# across them, so peak memory tracks the budget rather than the shard count. A per-file cap put the +# knob on the wrong axis: at 1000 rows each, resharding a dataset from 100 files to 10,000 took peak +# heap from 135 MB to 13.5 GB while describing exactly the same data. Ten thousand rows is ample for +# the statistics computed here (length quantiles, rates, cardinality); pass ``row_budget=None`` for a +# genuinely exhaustive scan. +DEFAULT_ROW_BUDGET = 10_000 + +# Rows read from a file however thin the budget gets. Below this a file cannot contribute the columns +# it alone witnesses, which is the whole reason every file is opened rather than a subset sampled. It +# is what makes the budget a target rather than a ceiling: 10,000 shards read this many each. +MIN_ROWS_PER_FILE = 10 def _format_of(path: str) -> str: @@ -66,14 +72,14 @@ def profile( source: FileSource, *, created_at: datetime | None = None, - row_cap: int | None = DEFAULT_ROW_CAP, + row_budget: int | None = DEFAULT_ROW_BUDGET, column_roles: dict[str, str] | None = None, ) -> DatasetProfile: """Profile the dataset behind ``source`` into a ``DatasetProfile``. - ``row_cap`` bounds how many rows are read from each file; ``None`` reads every row, which is - exact but scales memory with the dataset rather than the file count. Files smaller than the cap - are still read to the end, so a capped profile of a small dataset stays exhaustive. + ``row_budget`` bounds how many rows each *partition* reads in total, divided across its files; + ``None`` reads every row, which is exact but scales memory with the dataset. Files smaller than + their share are read to the end, so a budgeted profile of a small dataset stays complete. ``column_roles`` maps a column name to a role the caller is asserting, for datasets whose column names the role table does not recognize. Hints take precedence over name detection but still have @@ -110,7 +116,7 @@ def profile( groups = group_partitions(data_entries) for source_dir, partition_entries in groups: name = _partition_label(source_dir, len(groups)) - outcome = _profile_partition(source, name, source_dir, partition_entries, row_cap, column_roles or {}) + outcome = _profile_partition(source, name, source_dir, partition_entries, row_budget, column_roles or {}) partitions.append(outcome.partition) rows_scanned += outcome.rows_scanned files_read += outcome.files_read @@ -128,7 +134,7 @@ def profile( # Every data file, readable or not: the denominator that makes `files_read` a fraction rather # than a bare count. Non-data files (a README, a LICENSE) are not data and are counted nowhere. files_present=len(data_entries) + len(unreadable_files), - per_file_row_cap=row_cap, + row_budget=row_budget, seed=None, # head sampling makes no random choices; a seed would be theatre ) return DatasetProfile( @@ -140,6 +146,22 @@ def profile( ) +def _per_file_cap(row_budget: int | None, file_count: int) -> int | None: + """Split a partition's row budget across its files. + + Bounded below by :data:`MIN_ROWS_PER_FILE`, which is what makes the budget a target rather than a + ceiling: at a thousand shards the arithmetic share is ten rows, and at ten thousand it would be + one, which is too thin to witness a column. Overshooting the budget there is the right trade -- + the alternative is sampling a *subset of files*, which hides columns that appear only in later + shards, and file-level sampling is the tier of this problem still to solve. + """ + if row_budget is None: + return None + if file_count <= 1: + return row_budget + return max(MIN_ROWS_PER_FILE, row_budget // file_count) + + def _add_known(total: int | None, addend: int | None) -> int | None: """Sum two counts, where ``None`` means unknown and poisons the total. @@ -239,7 +261,7 @@ def _profile_partition( name: str, source_dir: str | None, entries: list[FileEntry], - row_cap: int | None, + row_budget: int | None, column_roles: dict[str, str], ) -> _PartitionOutcome: """Profile one partition — the files of one source directory, whatever formats they are in. @@ -261,6 +283,7 @@ def _profile_partition( files_read = 0 rows_present: int | None = 0 partition_scanned = True + row_cap = _per_file_cap(row_budget, len(entries)) read_strategy = "full" if row_cap is None else "head" split_profiles: list[SplitProfile] = [] for split in resolve_splits(entries): @@ -304,6 +327,7 @@ def _profile_partition( checksum=entry.checksum, file_format=file_format, read_strategy=read_strategy, + row_cap=row_cap, num_rows=num_rows, error=error, ) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index c9ff6f4acf..0e5f8d7b01 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -138,7 +138,9 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): # levels where each is decided, per file and per partition. assert {f.read_strategy for s in partition.splits for f in s.files} == {"head"} assert partition.stats_complete is True - assert result.sampling.per_file_row_cap == 1000 + assert result.sampling.row_budget == 10_000 + # The budget is split across the partition's two files, and each file records its own share. + assert {f.row_cap for s in partition.splits for f in s.files} == {5_000} assert result.sampling.rows_scanned == 3 assert result.sampling.rows_present == 3 assert result.sampling.files_read == result.sampling.files_present == 2 @@ -304,13 +306,13 @@ def test_profile_isolates_unreadable_files(tmp_path): assert result.sampling.files_present == 2 # ...out of two that were there to read -def test_profile_row_cap_bounds_reads_and_says_so(tmp_path): +def test_profile_row_budget_bounds_reads_and_says_so(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": i} for i in range(10)]) - result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=4) + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=4) assert result.sampling.rows_scanned == 4 - assert result.sampling.per_file_row_cap == 4 + assert result.sampling.row_budget == 4 assert result.partitions[0].stats_complete is False # 4 of 10 rows is not a full scan # The footer knows the total even though the cap stopped the read. Gating this on completeness # nulled it exactly when it carried information: "4 of 10" is a ratio, "4 of unknown" is not. @@ -321,10 +323,11 @@ def test_profile_row_cap_bounds_reads_and_says_so(tmp_path): def test_profile_uncapped_read_is_a_full_scan(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": i} for i in range(10)]) - result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=None) + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=None) assert {f.read_strategy for s in result.partitions[0].splits for f in s.files} == {"full"} - assert result.sampling.per_file_row_cap is None + assert result.sampling.row_budget is None + assert {f.row_cap for s in result.partitions[0].splits for f in s.files} == {None} assert result.partitions[0].stats_complete is True assert result.sampling.rows_scanned == result.sampling.rows_present == 10 @@ -334,7 +337,7 @@ def test_profile_cap_larger_than_a_jsonl_file_keeps_it_exhaustive(tmp_path): # Reading to EOF under the cap must stay exact, or capping would degrade every small dataset. (tmp_path / "train.jsonl").write_text('{"a": 1}\n{"a": 2}\n') - result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_cap=1000) + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=1000) assert result.partitions[0].splits[0].num_examples == 2 assert result.partitions[0].stats_complete is True @@ -571,3 +574,39 @@ def test_dataset_wide_completeness_is_one_expression(tmp_path): assert all(p.stats_complete for p in with_csv.partitions) # the parquet rows are still complete assert with_csv.unreadable_files # but there is data here that went unprofiled assert with_csv.sampling.files_read == 1 and with_csv.sampling.files_present == 2 + + +def test_row_budget_is_divided_across_a_partitions_files(tmp_path): + for shard in range(4): + _write_parquet(tmp_path / f"train-{shard:05d}-of-00004.parquet", [{"a": i} for i in range(200)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=400) + + assert {f.row_cap for s in result.partitions[0].splits for f in s.files} == {100} # 400 / 4 files + assert result.sampling.rows_scanned == 400 + assert result.sampling.row_budget == 400 + + +def test_rows_read_do_not_grow_when_a_dataset_is_resharded(tmp_path_factory): + # The property the budget exists for. Under a per-file cap the same data split ten ways further + # cost ten times the peak memory while describing exactly the same rows. + def rows_read(shards, per_shard): + root = tmp_path_factory.mktemp(f"shards{shards}") + for shard in range(shards): + _write_parquet(root / f"train-{shard:05d}-of-{shards:05d}.parquet", [{"a": i} for i in range(per_shard)]) + return profile(LocalFileSource(root), created_at=FIXED_TIME, row_budget=400).sampling.rows_scanned + + assert rows_read(4, 200) == rows_read(40, 20) == 400 + + +def test_row_budget_keeps_a_floor_under_very_thin_shards(tmp_path): + # Below the floor a file cannot witness the columns only it holds, which is the reason every file + # is opened rather than a subset sampled. Overshooting the budget there is the right trade, and + # the profile says so: row_cap is 10, not the 1 the arithmetic asked for. + for shard in range(10): + _write_parquet(tmp_path / f"train-{shard:05d}-of-00010.parquet", [{"a": i} for i in range(50)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=10) + + assert {f.row_cap for s in result.partitions[0].splits for f in s.files} == {10} + assert result.sampling.rows_scanned == 100 # deliberately over the budget From 8f5f858b4148851b047bc64dbd751fbc3224e7a5 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Wed, 5 Aug 2026 12:24:49 -0400 Subject: [PATCH 26/44] refactor(datasets): run the profiler as a job task, not a `nemo` CLI command `nemo datasets profile` was a published interface for a feature whose inputs and output contract have both moved several times this week. A subcommand is a promise to keep them still, and the profiler is not ready to make it. The audience today is the platform and the tests, neither of which needs a typed command. So the entry point becomes `python -m nemo_datasets_plugin.tasks.profile`: it reads a step config, profiles what the config names, and publishes the DatasetProfile as a job result artifact. The plugin now contributes no `nemo.cli` entry point at all. Landing the task here rather than with the Files integration keeps that branch to platform plumbing -- the job class, the API surface, the authz -- instead of also carrying profiler semantics. The seam falls where the dependencies already do: job results, step config and the sdk provider all exist on this branch, while the Files profile endpoint and the ranged-read fileset source do not. So the task profiles a local directory today, and the integration changes _build_source and adds the store step. The profiler core stays blind to where its bytes come from, which is what the FileSource seam is for. The step config also becomes the hint channel that --column-role was: `column_roles` is passed straight through, and a hint the data cannot support is reported as evidence on the profile rather than failing the task before it produces anything. typer and pyyaml were CLI-only and are dropped; nemo-platform-sdk is now a direct dependency because the task imports NeMoPlatform. pyarrow's floor moves 17 -> 19.0.1 to match every other declaration in the workspace. The uv.lock diff is larger than that implies: 333 of its 338 deleted lines are wheel URLs for armv7l, ppc64le, s390x and riscv64, none of which are in this workspace's platform markers. That is pre-existing staleness, not a consequence of these deps -- `uv lock` is a no-op on a clean tree but re-resolves and prunes on any edit to this plugin's dependency list, and every variant tried produced the same ~334 deletions. Signed-off-by: Albert Cui --- plugins/nemo-datasets/pyproject.toml | 11 +- .../src/nemo_datasets_plugin/cli.py | 80 -------- .../tasks/profile/__main__.py | 23 +++ .../nemo_datasets_plugin/tasks/profile/run.py | 173 ++++++++++++++++++ plugins/nemo-datasets/tests/test_cli.py | 79 -------- .../nemo-datasets/tests/test_profile_task.py | 144 +++++++++++++++ uv.lock | 8 +- 7 files changed, 349 insertions(+), 169 deletions(-) delete mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py create mode 100644 plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py delete mode 100644 plugins/nemo-datasets/tests/test_cli.py create mode 100644 plugins/nemo-datasets/tests/test_profile_task.py diff --git a/plugins/nemo-datasets/pyproject.toml b/plugins/nemo-datasets/pyproject.toml index 4a803821fd..0ba323e177 100644 --- a/plugins/nemo-datasets/pyproject.toml +++ b/plugins/nemo-datasets/pyproject.toml @@ -4,18 +4,19 @@ description = "Dataset profiler for NeMo Platform filesets." requires-python = ">=3.11,<3.15" dependencies = [ "nemo-platform-plugin", - "pyarrow>=17.0.0", - "pyyaml>=6.0.2", + "nemo-platform-sdk", + "pyarrow>=19.0.1", "pydantic>=2.10.3", - "typer>=0.20.0,<0.26", ] version = "0.1.0" [tool.uv.sources] nemo-platform-plugin = { workspace = true } +nemo-platform-sdk = { workspace = true } -[project.entry-points."nemo.cli"] -datasets = "nemo_datasets_plugin.cli:DatasetsCLI" +# Deliberately contributes no `nemo.cli` entry point. The profiler runs as a job task +# (`python -m nemo_datasets_plugin.tasks.profile`), which is invoked by the platform rather than +# typed by a user, so its inputs can keep moving while the feature is new. [build-system] requires = ["hatchling"] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py deleted file mode 100644 index 923cded46b..0000000000 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/cli.py +++ /dev/null @@ -1,80 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""The ``nemo datasets`` CLI — registered under the ``nemo.cli`` entry point.""" - -from __future__ import annotations - -import typer -from nemo_platform_plugin.cli import NemoCLI - - -class DatasetsCLI(NemoCLI): - """Exposes dataset commands as ``nemo datasets ...``.""" - - name = "datasets" - description = "Profile datasets stored as filesets." - - def get_cli(self) -> typer.Typer: - app = typer.Typer(help="Dataset profiling commands.") - - @app.callback() - def _root() -> None: - """Dataset profiling commands.""" - # A no-op callback keeps ``profile`` an explicit subcommand (``nemo datasets profile - # ``) instead of Typer collapsing the lone command into ``nemo datasets ``, - # which would break the moment a second command is added. - - @app.command() - def profile( - path: str = typer.Argument(..., help="Path to a local directory of dataset files."), - output: str = typer.Option("json", "--output", "-o", help="Output format: json | yaml."), - row_budget: int = typer.Option( - None, - "--row-budget", - help="Rows to read per partition, divided across its files (default 10000); 0 reads " - "every row, which is exact but scales memory with the dataset. A budget rather than a " - "per-file cap so peak memory does not grow when a dataset is resharded.", - min=0, - ), - column_role: list[str] = typer.Option( - None, - "--column-role", - help="Assert a column's role as NAME=ROLE (repeatable), e.g. --column-role q=prompt. " - "Takes precedence over name detection, but the dtype must still support the role; a " - "rejected hint is reported in the profile's evidence.", - metavar="NAME=ROLE", - ), - ) -> None: - """Profile a local dataset directory and print its DatasetProfile.""" - # Imported here, not at module scope: the platform calls get_cli() for every plugin at - # startup, and the profiler pulls in pyarrow. The budget default lives in the pipeline - # rather than being restated here, so an unspecified flag simply omits the argument. - from nemo_datasets_plugin.profiler.file_source import LocalFileSource - from nemo_datasets_plugin.profiler.pipeline import profile as run_profile - - if output not in {"json", "yaml"}: - raise typer.BadParameter("output must be 'json' or 'yaml'") - column_roles: dict[str, str] = {} - for pair in column_role or []: - name, separator, role = pair.partition("=") - if not separator or not name or not role: - raise typer.BadParameter(f"--column-role expects NAME=ROLE, got {pair!r}") - column_roles[name] = role - try: - source = LocalFileSource(path) - except NotADirectoryError as exc: - raise typer.BadParameter(str(exc)) from exc - - if row_budget is None: - result = run_profile(source, column_roles=column_roles) - else: - result = run_profile(source, row_budget=row_budget or None, column_roles=column_roles) - if output == "yaml": - import yaml - - typer.echo(yaml.safe_dump(result.model_dump(mode="json"), sort_keys=False)) - else: - typer.echo(result.model_dump_json(indent=2)) - - return app diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py new file mode 100644 index 0000000000..b54c16c6ef --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/__main__.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Module entry point: ``python -m nemo_datasets_plugin.tasks.profile``.""" + +import logging +import signal +import sys +from types import FrameType + +from nemo_datasets_plugin.tasks.profile.run import run + +logger = logging.getLogger(__name__) + + +def _shutdown_handler(signum: int, frame: FrameType | None) -> None: + logger.warning("Received shutdown signal (%s). Shutting down gracefully.", signum) + sys.exit(128 + signum) + + +if __name__ == "__main__": + signal.signal(signal.SIGTERM, _shutdown_handler) + sys.exit(run()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py new file mode 100644 index 0000000000..472091e179 --- /dev/null +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dataset-profiler task. + +Runs as a platform job: it reads a step config naming what to profile, runs the profiler, and +publishes the resulting ``DatasetProfile`` as a job result artifact named ``profile`` +(``profile.json``). + +This is deliberately *not* a ``nemo`` CLI subcommand. The profiler is new enough that its inputs and +its output contract are both still moving, and a published subcommand is a promise to keep them +still. A task module is invoked by the platform and by tests, which is the whole audience today. + +Only a local directory is profiled here. Reading a platform fileset through ranged requests, and +storing the profile back onto that fileset, both need Files-service surface this plugin does not +depend on; they arrive with the Files integration and change only :func:`_build_source` and the +publish step. The profiler core stays blind to where its bytes come from — that is what the +``FileSource`` seam is for. +""" + +from __future__ import annotations + +import json +import logging +import os +import tempfile +from pathlib import Path + +from nemo_datasets_plugin.profiler.file_source import FileSource, LocalFileSource +from nemo_datasets_plugin.profiler.pipeline import DEFAULT_ROW_BUDGET, profile +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.job_results import PlatformJobResults +from nemo_platform_plugin.jobs.constants import ( + EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, +) +from nemo_platform_plugin.sdk_provider import get_platform_sdk + +logger = logging.getLogger(__name__) + +# The service identity the task authenticates as. Any ``service:*`` principal is granted the internal +# ``ServiceSystem`` role, so no registration is required; a dedicated name just keeps audit logs and +# traces attributable. +_SERVICE_IDENTITY = "datasets" + +# Result artifact published back to the job's fileset. +_PROFILE_RESULT_NAME = "profile" +_PROFILE_FILE_NAME = "profile.json" + + +def run(sdk: NeMoPlatform | None = None) -> int: + """Entry point for the profiler task. Returns a process exit code.""" + _configure_logging() + try: + service_sdk = sdk or get_platform_sdk(as_service=_SERVICE_IDENTITY) + config = _load_step_config() + return _profile_and_publish( + service_sdk, + source=_build_source(config), + workspace=config.get("workspace") or _required_env(NEMO_JOB_WORKSPACE_ENVVAR), + job_name=_required_env(NEMO_JOB_ID_ENVVAR), + row_budget=_resolve_row_budget(config), + column_roles=_resolve_column_roles(config), + ) + except Exception: + logger.exception("Dataset profiler task failed") + return 1 + + +def _profile_and_publish( + sdk: NeMoPlatform, + *, + source: FileSource, + workspace: str, + job_name: str, + row_budget: int | None, + column_roles: dict[str, str], +) -> int: + logger.info("Profiling with a row budget of %s per partition", row_budget if row_budget else "unbounded") + dataset_profile = profile(source, row_budget=row_budget, column_roles=column_roles) + + # Scoped to the job's ephemeral storage when the platform provided one, and cleaned up either + # way — under the local subprocess backend this runs on a developer's machine, where an + # abandoned mkdtemp accumulates one directory per profiling run. + with tempfile.TemporaryDirectory( + prefix="dataset-profile-", + dir=os.environ.get(EPHEMERAL_TASK_STORAGE_PATH_ENVVAR) or None, + ) as scratch: + result_dir = Path(scratch) / _PROFILE_RESULT_NAME + result_dir.mkdir(parents=True) + (result_dir / _PROFILE_FILE_NAME).write_text(dataset_profile.model_dump_json(indent=2)) + + results = PlatformJobResults(job_name=job_name, workspace=workspace, sdk=sdk) + ref = results.save(_PROFILE_RESULT_NAME, result_dir) + logger.info("Published dataset profile: %s", ref.artifact_url) + return 0 + + +def _build_source(config: dict) -> FileSource: + """The files to profile, as named by the step config.""" + path = _required_config(config, "path") + try: + return LocalFileSource(path) + except NotADirectoryError as exc: + raise RuntimeError(f"step config 'path' must name a directory: {exc}") from exc + + +def _resolve_row_budget(config: dict) -> int | None: + """Rows the profiler may read per partition, from the step config. + + Defaults to the profiler's budget rather than an exhaustive read: uncapped, a partition holds + every row of every file in memory at roughly 20x the on-disk parquet size, which is what makes a + large fileset kill the job outright. A budgeted profile keeps exact row counts from the parquet + footers and reports ``stats_complete: false`` for the measurements, which is the trade the + sampling contract exists to describe. + + ``0`` asks for every row; use it when a proven value enumeration matters more than the cost. + """ + if "row_budget" not in config: + return DEFAULT_ROW_BUDGET + requested = config["row_budget"] + if requested is None: + return None + # Validated here as well as at any API boundary that produced it: this reads a file off disk, so + # nothing upstream is guaranteed to have checked it. + budget = int(requested) + if budget < 0: + raise ValueError(f"row_budget must be >= 0, got {budget}") + return budget or None + + +def _resolve_column_roles(config: dict) -> dict[str, str]: + """Caller-declared column roles, for datasets whose column names the role table does not know. + + Not validated against the role vocabulary here. The profiler applies its own dtype gates and + reports a hint the data cannot support as evidence on the profile, which is a better place for + the finding than a task that fails before producing anything. + """ + roles = config.get("column_roles") or {} + if not isinstance(roles, dict): + raise ValueError(f"column_roles must map column name to role, got {type(roles).__name__}") + return {str(name): str(role) for name, role in roles.items()} + + +def _load_step_config() -> dict: + path = os.environ.get(NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR) + if not path: + raise RuntimeError(f"{NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR} not set; running outside the platform?") + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + + +def _required_config(config: dict, key: str) -> str: + value = config.get(key) + if not value: + raise RuntimeError(f"Step config is missing '{key}'; nothing says what to profile.") + return str(value) + + +def _required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise RuntimeError(f"Missing required job environment variable: {name}") + return value + + +def _configure_logging() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) diff --git a/plugins/nemo-datasets/tests/test_cli.py b/plugins/nemo-datasets/tests/test_cli.py deleted file mode 100644 index 64d665cb6c..0000000000 --- a/plugins/nemo-datasets/tests/test_cli.py +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Smoke tests for the ``nemo datasets`` CLI surface.""" - -import json - -import pyarrow as pa -import pyarrow.parquet as pq -import typer -from nemo_datasets_plugin.cli import DatasetsCLI -from typer.testing import CliRunner - -runner = CliRunner() - - -def _mounted() -> typer.Typer: - """The app as the platform mounts it: ``nemo datasets ``.""" - root = typer.Typer() - root.add_typer(DatasetsCLI().get_cli(), name="datasets") - return root - - -def test_cli_metadata(): - cli = DatasetsCLI() - assert cli.name == "datasets" - assert cli.description - - -def test_profile_command_registered(): - app = DatasetsCLI().get_cli() - result = runner.invoke(app, ["--help"]) - assert result.exit_code == 0 - assert "profile" in result.stdout - - -def test_profile_command_profiles_a_directory(tmp_path): - pq.write_table( - pa.Table.from_pylist([{"prompt": "a"}, {"prompt": "b"}]), - tmp_path / "train-00000-of-00001.parquet", - ) - result = runner.invoke(_mounted(), ["datasets", "profile", str(tmp_path)]) - - assert result.exit_code == 0, result.output - payload = json.loads(result.stdout) - assert payload["partitions"][0]["features"][0]["name"] == "prompt" - assert payload["sampling"]["rows_scanned"] == 2 - - -def test_profile_command_rejects_unknown_output_format(tmp_path): - result = runner.invoke(_mounted(), ["datasets", "profile", str(tmp_path), "--output", "xml"]) - assert result.exit_code != 0 - - -def test_profile_command_rejects_non_directory(tmp_path): - target = tmp_path / "not-a-dir" - target.write_text("x") - result = runner.invoke(_mounted(), ["datasets", "profile", str(target)]) - assert result.exit_code != 0 - - -def test_profile_command_accepts_column_role_hints(tmp_path): - # The hint mechanism needs a caller on this branch; reading them from fileset metadata is the - # platform half and lands with the Files integration. - pq.write_table(pa.Table.from_pylist([{"q": "why?", "a": "because"}]), tmp_path / "train.parquet") - result = runner.invoke( - _mounted(), ["datasets", "profile", str(tmp_path), "--column-role", "q=prompt", "--column-role", "a=completion"] - ) - - assert result.exit_code == 0, result.output - payload = json.loads(result.stdout) - partition = payload["partitions"][0] - assert partition["classification"]["dataset_type"] == "prompt_completion" - assert [f["semantic_role_source"] for f in partition["features"]] == ["declared", "declared"] - - -def test_profile_command_rejects_a_malformed_column_role(tmp_path): - result = runner.invoke(_mounted(), ["datasets", "profile", str(tmp_path), "--column-role", "no-equals-sign"]) - assert result.exit_code != 0 diff --git a/plugins/nemo-datasets/tests/test_profile_task.py b/plugins/nemo-datasets/tests/test_profile_task.py new file mode 100644 index 0000000000..71d442e808 --- /dev/null +++ b/plugins/nemo-datasets/tests/test_profile_task.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the dataset-profiler job task.""" + +import json +from pathlib import Path +from typing import cast + +import nemo_datasets_plugin.tasks.profile.run as run_mod +import pyarrow as pa +import pyarrow.parquet as pq +from nemo_datasets_plugin.profiler.pipeline import DEFAULT_ROW_BUDGET +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.job_results import ResultRef +from nemo_platform_plugin.jobs.constants import ( + NEMO_JOB_ID_ENVVAR, + NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, + NEMO_JOB_WORKSPACE_ENVVAR, +) + +# The task touches the sdk only through PlatformJobResults, which is patched below, so a bare object +# stands in for it. +_SDK = cast(NeMoPlatform, object()) + + +def _dataset(root: Path, rows=None) -> Path: + root.mkdir(parents=True, exist_ok=True) + rows = rows or [{"q": "why?", "a": "because #### 4"}] + pq.write_table(pa.Table.from_pylist(rows), root / "train-00000-of-00001.parquet") + return root + + +def _install(monkeypatch, tmp_path: Path, config: dict) -> dict: + """Point the task at a step config and capture what it publishes.""" + published: dict = {} + + class _Results: + def __init__(self, *, job_name, workspace, sdk): + published.update(job_name=job_name, workspace=workspace) + + def save(self, name, local_path): + published["name"] = name + published["profile"] = json.loads((Path(local_path) / "profile.json").read_text()) + return ResultRef(name=name, artifact_url=f"file://{local_path}") + + config_path = tmp_path / "step-config.json" + config_path.write_text(json.dumps(config)) + monkeypatch.setenv(NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, str(config_path)) + monkeypatch.setenv(NEMO_JOB_WORKSPACE_ENVVAR, "ws1") + monkeypatch.setenv(NEMO_JOB_ID_ENVVAR, "job-1") + monkeypatch.setattr(run_mod, "PlatformJobResults", _Results) + return published + + +def test_task_profiles_a_directory_and_publishes_the_profile(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data)}) + + assert run_mod.run(_SDK) == 0 + + assert published["job_name"] == "job-1" + assert published["workspace"] == "ws1" + assert published["name"] == "profile" + profile = published["profile"] + assert profile["partitions"][0]["file_formats"] == ["parquet"] + assert profile["sampling"]["files_read"] == 1 + + +def test_task_prefers_the_step_configs_workspace_over_the_environment(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data), "workspace": "explicit"}) + + assert run_mod.run(_SDK) == 0 + assert published["workspace"] == "explicit" + + +def test_task_passes_column_role_hints_to_the_profiler(tmp_path, monkeypatch): + # The step config is the profiler's hint channel now that there is no CLI to carry --column-role. + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data), "column_roles": {"q": "prompt", "a": "completion"}}) + + assert run_mod.run(_SDK) == 0 + classification = published["profile"]["partitions"][0]["classification"] + assert classification["dataset_type"] == "prompt_completion" + + +def test_task_defaults_to_the_profilers_row_budget(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + published = _install(monkeypatch, tmp_path, {"path": str(data)}) + + assert run_mod.run(_SDK) == 0 + assert published["profile"]["sampling"]["row_budget"] == DEFAULT_ROW_BUDGET + + +def test_task_honours_an_explicit_row_budget(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data", rows=[{"a": i} for i in range(20)]) + published = _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": 5}) + + assert run_mod.run(_SDK) == 0 + assert published["profile"]["sampling"]["row_budget"] == 5 + assert published["profile"]["sampling"]["rows_scanned"] == 5 + + +def test_row_budget_zero_asks_for_every_row(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data", rows=[{"a": i} for i in range(20)]) + published = _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": 0}) + + assert run_mod.run(_SDK) == 0 + assert published["profile"]["sampling"]["row_budget"] is None + assert published["profile"]["partitions"][0]["stats_complete"] is True + + +def test_task_fails_when_the_step_config_says_nothing_to_profile(tmp_path, monkeypatch): + published = _install(monkeypatch, tmp_path, {}) + assert run_mod.run(_SDK) == 1 # a nonzero exit, not a traceback out of the container + assert published == {} + + +def test_task_fails_when_the_path_is_not_a_directory(tmp_path, monkeypatch): + target = tmp_path / "a-file" + target.write_text("x") + published = _install(monkeypatch, tmp_path, {"path": str(target)}) + + assert run_mod.run(_SDK) == 1 + assert published == {} + + +def test_task_fails_without_a_step_config(tmp_path, monkeypatch): + _install(monkeypatch, tmp_path, {"path": str(_dataset(tmp_path / "data"))}) + monkeypatch.delenv(NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR) + assert run_mod.run(_SDK) == 1 + + +def test_task_rejects_a_negative_row_budget(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": -1}) + assert run_mod.run(_SDK) == 1 + + +def test_task_rejects_column_roles_that_are_not_a_mapping(tmp_path, monkeypatch): + data = _dataset(tmp_path / "data") + _install(monkeypatch, tmp_path, {"path": str(data), "column_roles": ["q=prompt"]}) + assert run_mod.run(_SDK) == 1 diff --git a/uv.lock b/uv.lock index a80424c030..270c23e799 100644 --- a/uv.lock +++ b/uv.lock @@ -4404,19 +4404,17 @@ version = "0.1.0" source = { editable = "plugins/nemo-datasets" } dependencies = [ { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-platform-sdk", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pyarrow", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyyaml", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typer", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] requires-dist = [ { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, - { name = "pyarrow", specifier = ">=17.0.0" }, + { name = "nemo-platform-sdk", editable = "sdk/python/nemo-platform" }, + { name = "pyarrow", specifier = ">=19.0.1" }, { name = "pydantic", specifier = ">=2.10.3" }, - { name = "pyyaml", specifier = ">=6.0.2" }, - { name = "typer", specifier = ">=0.20.0,<0.26" }, ] [[package]] From 96cd5caebdb04d60b34e99780a5c5fd5bf8735ac Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 6 Aug 2026 12:28:16 -0400 Subject: [PATCH 27/44] docs(files): say why the profile contract lives in the shared package The module docstring justified its own location with a design that no longer exists: `DatasetMetadataContent` "can later carry it as a typed field". Storing a profile inside fileset metadata was abandoned precisely because writing one meant a read-modify-write of the whole metadata document, which clobbers any unrelated edit landing in between; a profile is its own entity now. That left the file looking parked under files/ by accident, which invites moving it into the datasets plugin. It should not move. The Files service is a first-class consumer -- its entities, endpoints, schemas and profile store all need this type -- so putting the contract in the plugin would make a core service depend on an optional one to deserialize rows in its own database. A deployment with no profiler installed still holds stored profiles and still answers GET .../filesets/{name}/profile. Records the real reason instead: pydantic-only and platform-free, so the profiler imports it standalone while Files imports it as the type it persists, and neither depends on the other. Also notes the `files/` placement is consistent with metadata.py, which houses the equally dataset-shaped DatasetMetadataContent. Docstring only; no behaviour change. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index a4c95463c4..a30e32b60b 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -19,8 +19,20 @@ tolerate unknown ones so the vocabulary can grow without a breaking change. Pydantic's default ``extra="ignore"`` gives the same forward-compatibility for unknown *fields*. -This module is pydantic-only — no platform dependencies — so the profiler can import it as a -standalone contract and ``DatasetMetadataContent`` can later carry it as a typed field. +**Why this lives in a shared package rather than with the profiler that writes it.** The Files +service is a first-class consumer, not a bystander: it stores a profile as its own entity and serves +it, so its entities, endpoints and schemas all need this type. Were the contract to live in the +datasets plugin, a core service would depend on an optional plugin to deserialize rows in its own +database — a deployment that installs no profiler still holds stored profiles and still answers +``GET .../filesets/{name}/profile``. Keeping it here means neither side depends on the other: the +module is pydantic-only, with no platform dependencies, so the profiler imports it standalone while +Files imports it as the type it persists. + +It sits under ``files/`` because Files is what stores and serves it, alongside the rest of that +service's shared contract — including ``metadata.py``, which houses the equally dataset-shaped +``DatasetMetadataContent``. Note that a profile is *not* carried inside fileset metadata: it is a +separate entity, so writing one cannot clobber an unrelated metadata edit that lands between a read +and a write. """ from __future__ import annotations From 4ff3bf204b5d178e7a2baef8d0e084fdda90fe56 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 6 Aug 2026 15:00:03 -0400 Subject: [PATCH 28/44] refactor(datasets): identify a partition by one name, not a label plus a key PartitionProfile carried both `name` (display) and `source_dir` (identity), and the label was worse than redundant. _partition_label returned "default" for a lone group *regardless of its directory*, so a dataset entirely under data/ reported name="default" alongside source_dir="data" -- the label did not summarize the identity, it discarded it. Its own description conceded the rest: "Display label, NOT a key. Derived from the layout, not guaranteed unique." A field that cannot be referenced, disagrees with the field that can, and is reproducible as `source_dir or "default"` is not earning a place in a stored contract. It only invited the question of which one to use. So there is one field. `name` is the path prefix a partition's files share -- a top-level directory, or "" at the fileset root. Unique by construction, because it is the grouping key itself rather than a projection of it. Empty is the one safe sentinel: no directory can be named it, so root-level files stay distinct from a directory literally called "default", which is the collision that made a label necessary in the first place. Keeping the field named `name` rather than introducing `id`: a partition's name being its identity is the ordinary expectation, and the anomaly was that it was not. It also survives the card-configs transition untouched -- HF calls it config_name, so a declared config populates this same field and the meaning holds. Visible change: a single-partition dataset under data/ now reports "data" instead of "default", and the nested data// layout likewise. That was a lie before. Display defaults move to the consumer, which is one expression: `name or "default"`. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 26 +++++-------- .../tests/files/test_dataset_profile.py | 13 +++---- .../profiler/partition.py | 30 +++++++-------- .../nemo_datasets_plugin/profiler/pipeline.py | 20 +--------- plugins/nemo-datasets/tests/test_pipeline.py | 38 +++++++++---------- 5 files changed, 50 insertions(+), 77 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index a30e32b60b..9d673397d2 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -406,27 +406,19 @@ class PartitionProfile(BaseModel): File membership and row counts live on ``splits`` — every file lands in exactly one split, so partition-level files / num_examples would be derivable duplication. - - ``source_dir`` is the identity and ``name`` is a label. They were a single string until the - consequences showed: root-level files and a directory literally named ``default`` collided under - one label, and dropping an unrelated file into a directory renamed that partition out of - existence — not changed, *gone*, so a stored reference resolved to nothing. """ name: str = Field( - default="default", - description=( - "Display label, NOT a key. Derived from the layout, not guaranteed unique, and free to " - "change when the layout does. Reference a partition by `source_dir` — or, once card " - "front-matter is parsed, by its declared config name." - ), - ) - source_dir: str | None = Field( - default=None, + default="", description=( - "Top-level directory whose files make up this partition; None when they sit at the " - "fileset root. The partition's identity: None and a directory named 'default' are " - 'different partitions even though both label as "default".' + "Identifies this partition, and unique within a profile. It is the path prefix its files " + 'share within the fileset: a top-level directory, or "" when they sit at the fileset ' + "root. Empty is a safe sentinel precisely because no directory can be named it, so " + "root-level files stay distinct from a directory literally called 'default'. Once card " + "front-matter is parsed, a declared config name populates this field instead — the same " + 'claim from a better source. For display, read it as `name or "default"`: storing that ' + "default was a lossy habit, because a lone partition under `data/` then reported " + '"default" and threw away the only thing identifying it.' ), ) file_formats: list[str] = Field( diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index 8915e435d1..668d7f642f 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -38,8 +38,7 @@ sampling: {rows_scanned: 2112, rows_present: 3201061, files_read: 33, files_present: 33, row_budget: 4096} partitions: - - name: default - source_dir: null + - name: "" file_formats: [parquet] stats_complete: false splits: @@ -82,8 +81,7 @@ sampling: {rows_scanned: 1024, rows_present: 46189, files_read: 2, files_present: 2, row_budget: 1024} partitions: - - name: default - source_dir: null + - name: "" file_formats: [parquet] stats_complete: false splits: @@ -126,8 +124,7 @@ sampling: {rows_scanned: 1024, rows_present: 21362, files_read: 2, files_present: 2, row_budget: 1024} partitions: - - name: default - source_dir: null + - name: "" file_formats: [parquet] stats_complete: false splits: @@ -189,7 +186,6 @@ def _build_profile() -> DatasetProfile: ), partitions=[ PartitionProfile( - source_dir=None, file_formats=["parquet"], stats_complete=False, splits=[ @@ -251,7 +247,8 @@ def test_fixture_deserializes(name): """Every fixture loads into the contract and round-trips.""" profile = DatasetProfile.model_validate(yaml.safe_load(FIXTURES[name])) assert profile.profile_schema_version == "1.0" - assert profile.partitions[0].name == "default" + # All three ship their shards at the fileset root, so the shared path prefix is empty. + assert profile.partitions[0].name == "" # Round-trip through JSON is lossless. assert DatasetProfile.model_validate_json(profile.model_dump_json()) == profile diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py index 0703c46cef..3216ee4c13 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/partition.py @@ -15,8 +15,11 @@ from nemo_datasets_plugin.profiler.splits import is_split_directory -def _top_dir(path: str) -> str | None: - """The partition directory for a file: its first path segment, else None for a root-level file. +def _top_dir(path: str) -> str: + """The partition a file belongs to: its first path segment, or ``""`` for a root-level file. + + Empty is a usable name precisely because no directory can be called it, so root-level files never + collide with a directory literally named ``default``. A split-named top-level directory (``train/``, ``test/``) is deliberately *not* a partition dimension. Grouping on it would split one dataset's train and test into unrelated partitions, @@ -25,29 +28,26 @@ def _top_dir(path: str) -> str | None: """ parts = PurePosixPath(path).parts if len(parts) <= 1 or is_split_directory(parts[0]): - return None + return "" return parts[0] -def group_partitions(entries: list[FileEntry]) -> list[tuple[str | None, list[FileEntry]]]: - """Group files into (source_dir, files) partitions by top-level directory. +def group_partitions(entries: list[FileEntry]) -> list[tuple[str, list[FileEntry]]]: + """Group files into (name, files) partitions by top-level directory, sorted by name. - Returns the *directory*, not a label — ``None`` for files at the fileset root. The label is the - caller's to derive, because the two are not the same thing: a lone group under ``data/`` labels - as "default" while its identity stays ``"data"``, and root-level files are a different partition - from a directory literally named ``default`` even though both label the same way. Collapsing - those into one string is what let a partition collide with another, and what let an unrelated - file rename one out of existence. + The name *is* the identity — the shared path prefix, not a display string derived from it. A lone + group under ``data/`` is named ``"data"``, not ``"default"``: reporting the latter discarded the + only thing identifying the partition, and left two partitions that could share a name. - Files whose top-level directory is a split name (``train/``, ``test/``) group under ``None`` + Files whose top-level directory is a split name (``train/``, ``test/``) group under ``""`` alongside root-level files: those are one dataset's splits, not separate partitions. """ - by_dir: dict[str | None, list[FileEntry]] = {} + by_dir: dict[str, list[FileEntry]] = {} for entry in entries: by_dir.setdefault(_top_dir(entry.path), []).append(entry) if len(by_dir) == 1: - # A single group is one partition holding everything, whatever its directory happened to be. + # A single group is one partition holding everything, keeping whatever directory it came from. return [(next(iter(by_dir)), list(entries))] - return sorted(by_dir.items(), key=lambda item: (item[0] is None, item[0] or "")) + return sorted(by_dir.items()) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 668dad7d3c..35e9b454e3 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -113,10 +113,8 @@ def profile( # None once any file's row count is unknown: the fileset's total is then unknowable, not zero. rows_present: int | None = 0 - groups = group_partitions(data_entries) - for source_dir, partition_entries in groups: - name = _partition_label(source_dir, len(groups)) - outcome = _profile_partition(source, name, source_dir, partition_entries, row_budget, column_roles or {}) + for name, partition_entries in group_partitions(data_entries): + outcome = _profile_partition(source, name, partition_entries, row_budget, column_roles or {}) partitions.append(outcome.partition) rows_scanned += outcome.rows_scanned files_read += outcome.files_read @@ -173,18 +171,6 @@ def _add_known(total: int | None, addend: int | None) -> int | None: return total + addend -def _partition_label(source_dir: str | None, group_count: int) -> str: - """The display label for a partition. Cosmetic only — ``source_dir`` carries the identity. - - A lone group labels as "default" whatever directory it came from, so the common - everything-under-``data/`` layout does not surface a meaningless container name. With several - groups each takes its directory name, and root-level files take "default". - """ - if group_count == 1 or source_dir is None: - return "default" - return source_dir - - def _unify_schemas(schemas: list[pa.Schema]) -> pa.Schema | None: """One schema describing every file of the partition, or None when they cannot be reconciled. @@ -259,7 +245,6 @@ class _PartitionOutcome: def _profile_partition( source: FileSource, name: str, - source_dir: str | None, entries: list[FileEntry], row_budget: int | None, column_roles: dict[str, str], @@ -353,7 +338,6 @@ def _profile_partition( ) partition = PartitionProfile( name=name, - source_dir=source_dir, # Summarized from the records rather than assumed: the partition no longer picks a format, # it reports the ones its files turned out to be in. file_formats=sorted({file.file_format for split in split_profiles for file in split.files if file.file_format}), diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 0e5f8d7b01..dedec09837 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -66,28 +66,30 @@ def test_resolve_splits_falls_back_to_single_default(): # --- partition grouping -------------------------------------------------------------------------- -def test_group_partitions_single_default_for_root_files(): +def test_group_partitions_names_the_root_partition_with_the_empty_prefix(): + # "" is the path prefix root-level files share, and no directory can be named it -- which is what + # keeps root files distinct from a directory literally called "default". assert group_partitions(_entries("train.parquet", "test.parquet")) == [ - (None, _entries("train.parquet", "test.parquet")) + ("", _entries("train.parquet", "test.parquet")) ] def test_group_partitions_collapses_single_container_dir(): - # One container directory is still one partition, but its identity stays the directory. Losing - # "data" here is what let a partition's identity move when the surrounding layout changed. + # One container directory is still one partition, and it keeps that directory as its name. + # Reporting "default" here discarded the only thing identifying the partition. parts = group_partitions(_entries("data/train.parquet", "data/test.parquet")) - assert [source_dir for source_dir, _ in parts] == ["data"] + assert [name for name, _ in parts] == ["data"] def test_group_partitions_splits_multiple_top_dirs(): parts = group_partitions(_entries("main/train.parquet", "socratic/train.parquet")) - assert [source_dir for source_dir, _ in parts] == ["main", "socratic"] + assert [name for name, _ in parts] == ["main", "socratic"] def test_group_partitions_does_not_treat_split_dirs_as_partitions(): # train/ and test/ are one dataset's splits, not two datasets. parts = group_partitions(_entries("train/data.parquet", "test/data.parquet")) - assert [source_dir for source_dir, _ in parts] == [None] + assert [name for name, _ in parts] == [""] def test_resolve_splits_reads_the_split_directory(): @@ -116,7 +118,7 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert result.profiler_info["name"] == "nemo-dataset-profiler" assert len(result.partitions) == 1 partition = result.partitions[0] - assert partition.name == "default" + assert partition.name == "" # root-level files: the empty path prefix assert partition.file_formats == ["parquet"] splits = {s.name: s for s in partition.splits} @@ -165,7 +167,6 @@ def test_profile_multiple_directories_become_partitions(tmp_path): assert [p.name for p in result.partitions] == ["main", "socratic"] assert all(p.file_formats == ["parquet"] for p in result.partitions) - assert [p.source_dir for p in result.partitions] == ["main", "socratic"] def test_profile_top_level_split_dirs_become_one_partition(tmp_path): @@ -176,7 +177,7 @@ def test_profile_top_level_split_dirs_become_one_partition(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert [p.name for p in result.partitions] == ["default"] + assert [p.name for p in result.partitions] == [""] splits = {s.name: s for s in result.partitions[0].splits} assert set(splits) == {"train", "test"} assert splits["train"].canonical == "train" @@ -191,7 +192,7 @@ def test_profile_nested_split_dirs_keep_splits_apart(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert [p.name for p in result.partitions] == ["default"] + assert [p.name for p in result.partitions] == ["data"] # the container directory, not "default" splits = {s.name: s for s in result.partitions[0].splits} assert set(splits) == {"train", "test"} assert splits["train"].num_examples == 2 @@ -209,8 +210,7 @@ def test_profile_keeps_a_mixed_format_directory_as_one_partition(tmp_path): assert len(result.partitions) == 1 partition = result.partitions[0] - assert partition.name == "default" - assert partition.source_dir == "data" + assert partition.name == "data" assert partition.file_formats == ["jsonl", "parquet"] assert {f.path.rsplit("/", 1)[-1]: f.file_format for s in partition.splits for f in s.files} == { "train-00000-of-00001.parquet": "parquet", @@ -224,10 +224,10 @@ def test_profile_keeps_a_mixed_format_directory_as_one_partition(tmp_path): def test_root_files_and_a_directory_named_default_stay_distinct(): - # Both label as "default"; only source_dir tells them apart. Flattening the two into one string - # produced two partitions with the same name and no way to reference either. + # The collision the empty-string sentinel exists to prevent: a derived label collapsed both to + # "default", leaving two partitions with one name and no way to reference either. parts = group_partitions(_entries("root.parquet", "default/inner.parquet")) - assert [source_dir for source_dir, _ in parts] == ["default", None] + assert [name for name, _ in parts] == ["", "default"] def test_an_unrelated_file_does_not_rename_a_partition(tmp_path): @@ -235,12 +235,12 @@ def test_an_unrelated_file_does_not_rename_a_partition(tmp_path): # renamed, *gone*, so a stored reference resolved to nothing. _write_parquet(tmp_path / "main" / "train.parquet", [{"q": "a"}]) _write_parquet(tmp_path / "socratic" / "train.parquet", [{"q": "b"}]) - before = [(p.name, p.source_dir) for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] + before = [p.name for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] (tmp_path / "main" / "notes.jsonl").write_text('{"note": "someone dropped this here"}\n') - after = [(p.name, p.source_dir) for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] + after = [p.name for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions] - assert before == after == [("main", "main"), ("socratic", "socratic")] + assert before == after == ["main", "socratic"] def test_profile_unions_columns_across_shards(tmp_path): From 61956c4682c2ce6ef12e42ac61d4514dca144c52 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 6 Aug 2026 17:25:38 -0400 Subject: [PATCH 29/44] refactor(datasets): enumerate the files that failed, count the ones that worked Measured on a profile whose consumer is meant to be an agent reading it into a context window: shards total ~tokens file records 8 3,810 952 1,064 (27%) 64 11,255 2,813 8,448 (75%) 512 70,332 17,583 67,072 (95%) At 512 shards, 95% of the payload was per-file records and every one of them said "this file was fine". Everything a reader actually wants -- schema, stats, classification, counts -- stayed flat at ~3.2k chars underneath. The profile was mostly a shard manifest wearing a dataset description. So SplitProfile.files becomes num_files, and FileRecord (8 fields) becomes FileError (path, error). Failures are named; successes are counted. Same profile now costs ~687 tokens at 8 shards or 512 -- flat, because the only unbounded part is gone. Nothing of value goes with it. `num_rows` already aggregated into SplitProfile.num_examples. `read_strategy` and `row_cap` are derivable from row_budget and num_files, both of which stay. `checksum` and `size_bytes` existed for the content digest, which was removed several commits ago and left them unread. `file_format` moves up to PartitionProfile.file_formats -- which stops being a denormalized summary and becomes the only place format is recorded, so its drift validator and the two tests guarding it go too. This also finishes a unification the coverage commit only half did. A .csv with no reader landed on the envelope while a corrupt parquet stayed buried in its split, though both are "a file I could not use". One list now, `DatasetProfile.file_errors`, sorted by path, whether or not a partition managed to group the file first. A reader asking "did anything go wrong?" reads one list whose length is the number of problems. Consequence worth stating plainly: there is now no per-file manifest at all, so the listing comparison I cited when removing the digest is no longer reconstructible either. `created_at` is the whole of what a profile says about its own freshness. That is deliberate while profiling is user-triggered and nothing consumes staleness; when something does, the cheap primitive is a backend version token, not a manifest rebuilt here. The contract says so rather than implying a mechanism that is gone. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 136 ++++++------------ .../tests/files/test_dataset_profile.py | 72 ++++------ .../nemo_datasets_plugin/profiler/pipeline.py | 64 ++++----- plugins/nemo-datasets/tests/test_pipeline.py | 73 ++++------ 4 files changed, 130 insertions(+), 215 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 9d673397d2..38746fed98 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -298,62 +298,21 @@ class ColumnStats(BaseModel): quality: TextQuality | None = Field(default=None, description="dtype == string: corruption signals") -class FileRecord(BaseModel): - """One physical file, measured. +class FileError(BaseModel): + """A file the profiler could not fully use, and why. - Stores the file's identity as the listing reported it — path, size, checksum — plus what the - reader learned cheaply. Concatenating ``files`` across a partition's splits reproduces that - partition's input list exactly, which is what lets a consumer compare a stored profile against a - fresh listing to see what changed. + Only failures are enumerated. Healthy files are counted (``SplitProfile.num_files``), because a + per-file record for each of them scaled the profile with shard count while telling a reader + nothing: at 512 shards those records were 95% of the payload and every one of them said "this + file was fine". Problems are the part worth naming, and there are few. """ path: str = Field(description="Relative path within the fileset.") - size_bytes: int - checksum: str | None = Field( - default=None, - description=( - 'As the Files service reports it (e.g. "sha256:..."), when it reports one at all — no backend ' - "does today. Without it, (path, size) is all there is to compare against a fresh listing: enough " - "to catch files added, removed, renamed or resized, but not a same-size in-place edit." - ), - ) - file_format: str | None = Field( - default=None, - description=( - "The format this file was read as (jsonl | parquet). A property of the file, not of the " - "partition holding it — which is why a partition may hold more than one. None only on a " - "profile written before formats were recorded per file." - ), - ) - read_strategy: str | None = Field( - default=None, - description=( - "How this file's rows were sampled: full | head. The *policy* applied, not the outcome — " - "a head-capped read of a file smaller than the cap still says head, and whether it ended " - "up complete is `num_rows` versus what was scanned. Per file because it follows format, " - "which is also per file: a parquet shard can be sampled by row group where a jsonl file " - "in the same partition can only be read from the top." - ), - ) - row_cap: int | None = Field( - default=None, + error: str = Field( description=( - "Rows this file's read was bounded to, None when unbounded. Derived from " - "`SamplingInfo.row_budget` divided across the partition's files, so it is per file rather " - "than a global setting: the same budget yields 1000 rows each across ten shards and ten " - "each across a thousand, which is what keeps peak memory flat as a dataset is resharded." - ), - ) - num_rows: int | None = Field( - default=None, - description="Exact only (parquet footer / exhaustive scan), else None.", - ) - error: str | None = Field( - default=None, - description=( - "Why this file was not fully read, when it wasn't — unreadable, corrupt, or partially " - "parsed. None means a clean read. Without it a missing `num_rows` is indistinguishable " - "from a profiler bug, and a consumer cannot tell corrupt input from unsupported input." + "Why this file was not fully read: unreadable, corrupt, partially parsed, or in a format " + "with no reader. A file that was read cleanly never appears here, so the absence of a path " + "is itself the claim that it was fine." ), ) @@ -383,11 +342,13 @@ class SplitProfile(BaseModel): "train, with the variant's intent kept in `name`." ), ) - files: list[FileRecord] = Field( + num_files: int = Field( + default=0, description=( - "Every file resolved into this split, measured. Partitioning is exhaustive and disjoint: each " - "file of the partition lands in exactly one split, so concatenating `files` across splits " - "reconstructs the partition's file list with no gaps or repeats." + "How many files resolved into this split. Partitioning is exhaustive and disjoint — each file " + "of the partition lands in exactly one split — so these sum to the partition's file count. " + "A count rather than a list: the paths of healthy shards are the one part of a profile that " + "grows without bound and informs no decision." ), ) num_examples: int | None = Field( @@ -424,11 +385,11 @@ class PartitionProfile(BaseModel): file_formats: list[str] = Field( default_factory=list, description=( - "The distinct formats among this partition's files, sorted — normally exactly one. " - "Format is a property of a file (see `FileRecord.file_format`), never a partition " - "dimension: a stray .jsonl beside .parquet shards is noise, not a second dataset, so it " - "stays in this partition and shows up here. jsonl | parquet are read today; csv | arrow " - "are reserved vocabulary the profiler cannot read yet and reports as unsupported." + "The distinct formats this partition's files are in, sorted — normally exactly one, and " + "more than one when a stray .jsonl sits beside .parquet shards. That is noise, not a " + "second dataset, so it stays in this partition and shows up here rather than splitting it. " + "jsonl | parquet are read today; csv | arrow are reserved vocabulary the profiler cannot " + "read yet and reports on `DatasetProfile.file_errors` instead." ), ) splits: list[SplitProfile] = Field(description="card-declared > path-detected > single 'default' split.") @@ -463,20 +424,6 @@ def _stats_keys_subset_of_features(self) -> PartitionProfile: raise ValueError(f"stats keys must name top-level features; unknown columns: {sorted(unknown)}") return self - @model_validator(mode="after") - def _file_formats_cover_the_records(self) -> PartitionProfile: - """Every format recorded on a file must appear in the partition's summary. - - A subset check rather than equality, so the two cannot drift in the direction that matters: - a summary omitting a format that demonstrably exists is wrong, while a profile written - before formats were recorded per file has nothing on its records and nothing to check. - """ - recorded = {file.file_format for split in self.splits for file in split.files if file.file_format} - missing = recorded - set(self.file_formats) - if missing: - raise ValueError(f"file_formats omits formats present on this partition's files: {sorted(missing)}") - return self - # ---- envelope ---------------------------------------------------------------------------------- @@ -493,7 +440,7 @@ class SamplingInfo(BaseModel): The dataset-wide question is still one expression away, and now says which half failed:: - all(p.stats_complete for p in profile.partitions) and not profile.unreadable_files + all(p.stats_complete for p in profile.partitions) and not profile.file_errors """ rows_scanned: int = Field(description="Total rows actually parsed across all files.") @@ -514,7 +461,7 @@ class SamplingInfo(BaseModel): description=( "Data files the fileset holds, whether or not this run could read them — the denominator " "`files_read` is a fraction of. Includes files in formats with no reader, since those are " - "data that went unprofiled (they are listed in `DatasetProfile.unreadable_files`). A README " + "data that went unprofiled (they are named on `DatasetProfile.file_errors`). A README " "is not data and is counted nowhere. Every readable file should be opened, since " "head-sampling a *subset of files* hides columns that appear only in later shards, so expect " "these two to match until scale forces file-level sampling." @@ -526,9 +473,9 @@ class SamplingInfo(BaseModel): "Rows the caller allowed per partition, None for an unbounded read. A budget rather than a " "per-file cap because the cost is per partition: a per-file cap made peak memory scale with " "shard count, so the same dataset resharded from 100 files to 10,000 went from megabytes to " - "gigabytes without holding any more data. The per-file cap this produced is on each " - "`FileRecord.row_cap`. Not a hard ceiling: every file is still read at least a few rows, " - "since a file sampled too thinly cannot contribute the columns it alone witnesses." + "gigabytes without holding any more data. Not a hard ceiling: every file is still read at " + "least a few rows, since one sampled too thinly cannot contribute the columns it alone " + "witnesses, so a partition with very many files may exceed its budget." ), ) seed: int | None = Field(default=None, description="RNG seed used for row selection, for reproducibility.") @@ -537,13 +484,16 @@ class SamplingInfo(BaseModel): class DatasetProfile(BaseModel): """The machine-owned dataset profile — the root of the stored contract. - Deliberately carries no staleness marker. A stored digest would freeze "which files count as - inputs" into the data at write time, and that judgment moves: once card front-matter drives - split declaration, ``README.md`` becomes an input. Changing the rule would then invalidate every - stored profile at once, with no way to tell a real change from a definition change. The - ``FileRecord``s already describe the inputs, so a consumer that needs to know whether a profile - is current compares them against a fresh listing — same cost, and it learns *what* changed - rather than merely *that* something did. + Deliberately carries no staleness marker, and no per-file manifest to reconstruct one from. A + stored digest would freeze "which files count as inputs" into the data at write time, and that + judgment moves: once card front-matter drives split declaration, ``README.md`` becomes an input. + Changing the rule would then invalidate every stored profile at once, with no way to tell a real + change from a definition change. + + So a profile says when it was made and nothing about whether it still holds. ``created_at`` is + the whole of it. That is deliberate while profiling is user-triggered and nothing consumes + freshness; when something does, the cheap primitive is a fileset version token from the storage + backend, which costs no listing and freezes no policy — not a manifest reconstructed here. """ profile_schema_version: str = Field( @@ -559,15 +509,15 @@ class DatasetProfile(BaseModel): partitions: list[PartitionProfile] = Field( description="Single partition in the common homogeneous case; there is no fileset-level rollup.", ) - unreadable_files: list[FileRecord] = Field( + file_errors: list[FileError] = Field( default_factory=list, description=( - "Files that plainly hold dataset records but that no partition could take, because the " - "profiler has no reader for their format; each carries the reason on `error`. Reporting " - "them is what keeps a directory of .csv shards from profiling as an exhaustively scanned " - "*empty* dataset, indistinguishable from one that really is empty. A file whose format is " - "known but whose read failed keeps its FileRecord inside its split instead — it was " - "grouped and attempted, these never were." + "Every file the profiler could not fully use, from anywhere in the fileset: a format with " + "no reader, a corrupt shard, a partially parsed one. Reporting them is what keeps a " + "directory of .csv shards from profiling as an exhaustively scanned *empty* dataset, " + "indistinguishable from one that really is empty. One list rather than two, because " + '"a file I could not use" is the same finding whether or not a partition managed to group ' + 'it first, and a reader asking "did anything go wrong?" should not have to look twice.' ), ) diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index 668d7f642f..b35db8ed60 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -19,7 +19,7 @@ DatasetProfile, Evidence, FeatureSchema, - FileRecord, + FileError, MessageStats, PartitionClassification, PartitionProfile, @@ -42,12 +42,8 @@ file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 3200861, - files: [{path: train-00000-of-00032.parquet, size_bytes: 193777041, - checksum: sha256:9c1e..., num_rows: 100027, file_format: parquet, read_strategy: head, row_cap: 512}]} - - {name: test, canonical: test, num_examples: 200, - files: [{path: test-00000-of-00001.parquet, size_bytes: 411552, - checksum: sha256:02af..., num_rows: 200, file_format: parquet, read_strategy: head, row_cap: 512}]} + - {name: train, canonical: train, num_examples: 3200861, num_files: 32} + - {name: test, canonical: test, num_examples: 200, num_files: 1} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -85,12 +81,8 @@ file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 43835, - files: [{path: train-00000-of-00001.parquet, size_bytes: 22105331, - checksum: sha256:77b0..., num_rows: 43835, file_format: parquet, read_strategy: head, row_cap: 512}]} - - {name: test, canonical: test, num_examples: 2354, - files: [{path: test-00000-of-00001.parquet, size_bytes: 1198422, - checksum: sha256:5c1d..., num_rows: 2354, file_format: parquet, read_strategy: head, row_cap: 512}]} + - {name: train, canonical: train, num_examples: 43835, num_files: 1} + - {name: test, canonical: test, num_examples: 2354, num_files: 1} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -128,12 +120,8 @@ file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 20324, - files: [{path: train-00000-of-00001.parquet, size_bytes: 44201991, - checksum: sha256:e410..., num_rows: 20324, file_format: parquet, read_strategy: head, row_cap: 512}]} - - {name: validation, canonical: validation, num_examples: 1038, - files: [{path: validation-00000-of-00001.parquet, size_bytes: 2311008, - checksum: sha256:8bd2..., num_rows: 1038, file_format: parquet, read_strategy: head, row_cap: 512}]} + - {name: train, canonical: train, num_examples: 20324, num_files: 1} + - {name: validation, canonical: validation, num_examples: 1038, num_files: 1} features: - {name: prompt, dtype: string, semantic_role: prompt, semantic_role_source: detected} - {name: response, dtype: string, semantic_role: completion, semantic_role_source: detected} @@ -193,17 +181,7 @@ def _build_profile() -> DatasetProfile: name="train", canonical="train", num_examples=2048, - files=[ - FileRecord( - path="train-00000.parquet", - size_bytes=123, - checksum="sha256:ab", - file_format="parquet", - read_strategy="head", - row_cap=512, - num_rows=2048, - ) - ], + num_files=1, ) ], features=[ @@ -357,24 +335,28 @@ def test_unknown_fields_are_ignored_for_forward_compat(): assert profile.partitions[0].classification.dataset_type == "scored_response" -def test_file_formats_must_not_omit_a_format_its_files_carry(): - # The partition summary is derived from the records, so the two can drift. A summary claiming - # one format while a file says otherwise would report the partition as homogeneous when it is - # not -- the very assumption that made format a partition dimension and cost names their - # stability. Subset, not equality, so a profile written before per-file formats still loads. +def test_file_errors_are_the_only_channel_for_trouble(): + # Healthy files are counted, never listed, so a reader asking "did anything go wrong?" reads one + # list whose length is the number of problems -- not one that grows with the shard count and is + # 95% success records at scale. doc = yaml.safe_load(HELPSTEER2) - doc["partitions"][0]["splits"][0]["files"][0]["file_format"] = "jsonl" - with pytest.raises(ValueError, match="file_formats omits"): - DatasetProfile.model_validate(doc) + doc["file_errors"] = [ + {"path": "train-00007-of-00032.parquet", "error": "ArrowInvalid: not a parquet file"}, + {"path": "notes.csv", "error": "no reader for '.csv' files"}, + ] + profile = DatasetProfile.model_validate(doc) + assert [e.path for e in profile.file_errors] == ["train-00007-of-00032.parquet", "notes.csv"] + # A shard the profiler could not read and a format it has no reader for are the same finding, + # and land in the same place whether or not a partition managed to group the file first. + assert all(isinstance(e, FileError) and e.error for e in profile.file_errors) + assert DatasetProfile.model_validate_json(profile.model_dump_json()) == profile -def test_a_partition_written_before_per_file_formats_still_loads(): - doc = yaml.safe_load(HELPSTEER2) - for split in doc["partitions"][0]["splits"]: - for file in split["files"]: - del file["file_format"] - profile = DatasetProfile.model_validate(doc) - assert profile.partitions[0].splits[0].files[0].file_format is None + +def test_a_clean_profile_names_no_files_at_all(): + profile = DatasetProfile.model_validate(yaml.safe_load(HELPSTEER2)) + assert profile.file_errors == [] + assert [s.num_files for s in profile.partitions[0].splits] == [1, 1] def test_a_profile_written_before_the_digest_was_dropped_still_loads(): diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 35e9b454e3..20455e1419 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -34,7 +34,7 @@ DatasetProfile, Evidence, FeatureSchema, - FileRecord, + FileError, PartitionClassification, PartitionProfile, SamplingInfo, @@ -96,11 +96,9 @@ def profile( # scanned, empty dataset — indistinguishable from a dataset that really is empty. They get real # FileRecords like any other file the profiler could not read, just at the envelope, since no # partition ever grouped them. - unreadable_files = [ - FileRecord( + file_errors = [ + FileError( path=entry.path, - size_bytes=entry.size_bytes, - checksum=entry.checksum, error=f"no reader for '{PurePosixPath(entry.path).suffix.lower()}' files", ) for entry in sorted(all_entries, key=lambda entry: entry.path) @@ -119,10 +117,11 @@ def profile( rows_scanned += outcome.rows_scanned files_read += outcome.files_read rows_present = _add_known(rows_present, outcome.rows_present) + file_errors.extend(outcome.file_errors) - # A format with no reader holds an unknown number of rows, so it makes the fileset total unknown - # in exactly the way an unread file does. - if unreadable_files: + # A file the profiler could not use holds an unknown number of rows, so it makes the fileset + # total unknown — whether it was skipped for want of a reader or failed mid-read. + if file_errors: rows_present = None sampling = SamplingInfo( @@ -131,7 +130,7 @@ def profile( files_read=files_read, # files actually opened and read, not files merely listed # Every data file, readable or not: the denominator that makes `files_read` a fraction rather # than a bare count. Non-data files (a README, a LICENSE) are not data and are counted nowhere. - files_present=len(data_entries) + len(unreadable_files), + files_present=len(data_entries) + sum(1 for e in file_errors if detect_format(e.path) is None), row_budget=row_budget, seed=None, # head sampling makes no random choices; a seed would be theatre ) @@ -140,7 +139,9 @@ def profile( profiler_info={"name": PROFILER_NAME, "version": PROFILER_VERSION}, sampling=sampling, partitions=partitions, - unreadable_files=unreadable_files, + # Sorted so a reader scanning for trouble sees it in a stable order, whatever partition it + # came from; partitions contribute theirs as they are profiled. + file_errors=sorted(file_errors, key=lambda error: error.path), ) @@ -240,6 +241,7 @@ class _PartitionOutcome: rows_scanned: int files_read: int # files actually opened and read, so `files_read` can exclude failures rows_present: int | None # rows known to exist here, or None once any file's count is unknown + file_errors: list[FileError] # files this partition grouped but could not fully read def _profile_partition( @@ -257,9 +259,9 @@ def _profile_partition( instead flow through to ``_measure``, which infers the schema from rows when not every file declared one. - An unreadable file (or a format with no registered reader) is isolated: it keeps its FileRecord, - records *why* on ``FileRecord.error``, contributes no rows, and flips ``scanned_all`` off — it - never aborts the profile. + An unreadable file (or a format with no registered reader) is isolated: it is named on a + :class:`FileError` the envelope collects, contributes no rows, and flips ``scanned_all`` off — it + never aborts the profile. Files that read cleanly are counted, not listed. """ partition_rows: list[dict] = [] arrow_schemas: list[pa.Schema] = [] @@ -269,18 +271,18 @@ def _profile_partition( rows_present: int | None = 0 partition_scanned = True row_cap = _per_file_cap(row_budget, len(entries)) - read_strategy = "full" if row_cap is None else "head" + file_errors: list[FileError] = [] + file_formats: set[str] = set() split_profiles: list[SplitProfile] = [] for split in resolve_splits(entries): - file_records: list[FileRecord] = [] split_examples = 0 split_counts_known = True # every file's exact total row count is known (footer or full scan) split_scanned = True # every row of every file was actually parsed for entry in split.entries: - file_format = _format_of(entry.path) + file_formats.add(_format_of(entry.path)) error: str | None = None try: - result = get_reader(file_format).read(source, entry, row_cap=row_cap) + result = get_reader(_format_of(entry.path)).read(source, entry, row_cap=row_cap) except Exception as exc: # Failure isolation: an unreadable file (or missing reader) keeps its identity, # skips its rows, and does not abort the profile. The reason is recorded rather than @@ -305,18 +307,8 @@ def _profile_partition( # Exhaustive requires parsing every row; a known footer count alone is not enough, and # a partial read (corrupt lines skipped) is not exhaustive however many rows it got. scanned_all = num_rows is not None and result.rows_scanned >= num_rows and error is None - file_records.append( - FileRecord( - path=entry.path, - size_bytes=entry.size_bytes, - checksum=entry.checksum, - file_format=file_format, - read_strategy=read_strategy, - row_cap=row_cap, - num_rows=num_rows, - error=error, - ) - ) + if error is not None: + file_errors.append(FileError(path=entry.path, error=error)) rows_present = _add_known(rows_present, num_rows) if num_rows is None: split_counts_known = False @@ -329,7 +321,7 @@ def _profile_partition( SplitProfile( name=split.name, canonical=split.canonical, - files=file_records, + num_files=len(split.entries), num_examples=split_examples if split_counts_known else None, ) ) @@ -338,9 +330,9 @@ def _profile_partition( ) partition = PartitionProfile( name=name, - # Summarized from the records rather than assumed: the partition no longer picks a format, - # it reports the ones its files turned out to be in. - file_formats=sorted({file.file_format for split in split_profiles for file in split.files if file.file_format}), + # Observed, not chosen: the partition reports the formats its files turned out to be in + # rather than picking one and splitting to keep that true. + file_formats=sorted(file_formats), splits=split_profiles, features=features, stats=stats, @@ -350,5 +342,9 @@ def _profile_partition( classification=classification, ) return _PartitionOutcome( - partition=partition, rows_scanned=rows_scanned, files_read=files_read, rows_present=rows_present + partition=partition, + rows_scanned=rows_scanned, + files_read=files_read, + rows_present=rows_present, + file_errors=file_errors, ) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index dedec09837..2597f86d9a 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -126,7 +126,7 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert splits["train"].canonical == "train" assert splits["train"].num_examples == 2 assert splits["validation"].num_examples == 1 - assert splits["train"].files[0].num_rows == 2 + assert splits["train"].num_files == 1 # Row schema, stats, and classification are all derived now. assert [f.name for f in partition.features] == ["prompt"] @@ -135,14 +135,10 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): assert partition.stats["prompt"].text is not None assert partition.classification.dataset_type == "prompt_only" # a lone prompt column, no target - # read_strategy is the policy, stats_complete is the outcome: a capped run over files that all - # fit under the cap is still a complete scan, which is why the two live apart -- and now at the - # levels where each is decided, per file and per partition. - assert {f.read_strategy for s in partition.splits for f in s.files} == {"head"} + # A budgeted run over files that all fit under their share is still a complete scan, which is + # why the budget and the outcome are separate fields. assert partition.stats_complete is True assert result.sampling.row_budget == 10_000 - # The budget is split across the partition's two files, and each file records its own share. - assert {f.row_cap for s in partition.splits for f in s.files} == {5_000} assert result.sampling.rows_scanned == 3 assert result.sampling.rows_present == 3 assert result.sampling.files_read == result.sampling.files_present == 2 @@ -212,10 +208,6 @@ def test_profile_keeps_a_mixed_format_directory_as_one_partition(tmp_path): partition = result.partitions[0] assert partition.name == "data" assert partition.file_formats == ["jsonl", "parquet"] - assert {f.path.rsplit("/", 1)[-1]: f.file_format for s in partition.splits for f in s.files} == { - "train-00000-of-00001.parquet": "parquet", - "extra.jsonl": "jsonl", - } # Both formats' columns reach features. Trusting the declared parquet schema would have erased # `question`, which only the schemaless file witnesses -- the defect the split worked around. assert sorted(f.name for f in partition.features) == ["prompt", "question"] @@ -298,8 +290,8 @@ def test_profile_isolates_unreadable_files(tmp_path): splits = {s.name: s for s in result.partitions[0].splits} assert splits["train"].num_examples == 1 assert splits["test"].num_examples is None # unreadable -> count unknown, not a crash - assert splits["test"].files[0].num_rows is None - assert splits["test"].files[0].error is not None # ...and the profile says why + assert [e.path for e in result.file_errors] == ["test-00000-of-00001.parquet"] # named, with a reason + assert result.file_errors[0].error assert result.partitions[0].stats_complete is False # a file could not be fully parsed assert result.sampling.rows_present is None assert result.sampling.files_read == 1 # one file was actually read; the other never opened @@ -325,9 +317,7 @@ def test_profile_uncapped_read_is_a_full_scan(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=None) - assert {f.read_strategy for s in result.partitions[0].splits for f in s.files} == {"full"} assert result.sampling.row_budget is None - assert {f.row_cap for s in result.partitions[0].splits for f in s.files} == {None} assert result.partitions[0].stats_complete is True assert result.sampling.rows_scanned == result.sampling.rows_present == 10 @@ -357,8 +347,8 @@ def test_profile_reports_unsupported_data_files(tmp_path): assert result.sampling.files_read == 0 assert result.sampling.files_present == 2 # both are data; neither could be read # Typed records now, each saying why -- not bare paths tucked into a free-form dict. - assert [f.path for f in result.unreadable_files] == ["test.arrow", "train.csv"] - assert all("no reader" in f.error for f in result.unreadable_files) + assert [e.path for e in result.file_errors] == ["test.arrow", "train.csv"] + assert all("no reader" in e.error for e in result.file_errors) def test_profile_ignores_non_data_files_without_penalty(tmp_path): @@ -370,7 +360,7 @@ def test_profile_ignores_non_data_files_without_penalty(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) assert result.partitions[0].stats_complete is True - assert result.unreadable_files == [] + assert result.file_errors == [] assert result.sampling.files_present == 1 # the README and LICENSE are not data, counted nowhere @@ -381,9 +371,9 @@ def test_profile_records_a_partial_jsonl_read(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - record = result.partitions[0].splits[0].files[0] - assert record.num_rows == 2 # the readable rows survived - assert record.error is not None and "line 2" in record.error + assert result.partitions[0].splits[0].num_examples == 2 # the readable rows survived + assert [e.path for e in result.file_errors] == ["train.jsonl"] + assert "line 2" in result.file_errors[0].error assert result.partitions[0].stats_complete is False # a line was lost, so not a full scan @@ -472,11 +462,13 @@ def test_profile_survives_a_hostile_directory(tmp_path): # Nothing here is exhaustive, and the profile says so rather than looking clean. assert result.partitions[0].stats_complete is False assert result.sampling.rows_present is None - assert [f.path for f in result.unreadable_files] == ["leftovers.csv"] - - records = {f.path: f for p in result.partitions for s in p.splits for f in s.files} - assert records["train-00001-of-00002.parquet"].error is not None # corrupt file, named and explained - assert records["extra.jsonl"].error is not None # partial parse, named and explained + # One channel for every file the profiler could not use, whether or not a partition grouped it: + # the .csv it never read, the corrupt shard, and the jsonl it only partly parsed. + assert [e.path for e in result.file_errors] == [ + "extra.jsonl", + "leftovers.csv", + "train-00001-of-00002.parquet", + ] # One partition, not one per format: the stray .jsonl is noise, not a second dataset. assert len(result.partitions) == 1 @@ -491,11 +483,10 @@ def test_profile_survives_a_hostile_directory(tmp_path): assert DatasetProfile.model_validate_json(result.model_dump_json()) == result -def test_stored_file_records_reproduce_the_input_list(tmp_path): - # The contract promises split membership is exhaustive and disjoint, which is what lets a - # consumer compare a stored profile against a fresh listing to decide whether it is current. - # That comparison is the whole reason the records carry path/size/checksum, so the invariant is - # worth asserting directly rather than through a digest that happened to depend on it. +def test_split_file_counts_account_for_every_data_file(tmp_path): + # The contract promises split membership is exhaustive and disjoint. With per-file records gone + # the counts are all that carries it, so the invariant is worth asserting on them directly -- + # a count that silently dropped a file would look exactly like a smaller dataset. _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"a": 1}]) _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"a": 2}]) _write_parquet(tmp_path / "test-00000-of-00001.parquet", [{"a": 3}]) @@ -504,10 +495,9 @@ def test_stored_file_records_reproduce_the_input_list(tmp_path): source = LocalFileSource(tmp_path) result = profile(source, created_at=FIXED_TIME) - stored = [f.path for partition in result.partitions for split in partition.splits for f in split.files] + counted = sum(split.num_files for partition in result.partitions for split in partition.splits) listed = [e.path for e in source.list_files() if e.path.endswith(".parquet")] - assert sorted(stored) == sorted(listed) # exhaustive - assert len(stored) == len(set(stored)) # and disjoint + assert counted == len(listed) # exhaustive and disjoint: each file lands in exactly one split def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): @@ -521,8 +511,7 @@ def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise - records = {f.path: f for p in result.partitions for s in p.splits for f in s.files} - assert records["extra.xyz"].num_rows is None # kept, but unreadable + assert "extra.xyz" in {e.path for e in result.file_errors} # named, not silently dropped assert result.partitions[0].stats_complete is False @@ -567,12 +556,12 @@ def test_dataset_wide_completeness_is_one_expression(tmp_path): # says *which* half failed, which the single bit could not. _write_parquet(tmp_path / "train.parquet", [{"a": 1}]) clean = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert all(p.stats_complete for p in clean.partitions) and not clean.unreadable_files + assert all(p.stats_complete for p in clean.partitions) and not clean.file_errors (tmp_path / "extra.csv").write_text("a,b\n1,2\n") with_csv = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) assert all(p.stats_complete for p in with_csv.partitions) # the parquet rows are still complete - assert with_csv.unreadable_files # but there is data here that went unprofiled + assert with_csv.file_errors # but there is data here that went unprofiled assert with_csv.sampling.files_read == 1 and with_csv.sampling.files_present == 2 @@ -582,8 +571,7 @@ def test_row_budget_is_divided_across_a_partitions_files(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=400) - assert {f.row_cap for s in result.partitions[0].splits for f in s.files} == {100} # 400 / 4 files - assert result.sampling.rows_scanned == 400 + assert result.sampling.rows_scanned == 400 # 400 / 4 files = 100 rows each assert result.sampling.row_budget == 400 @@ -602,11 +590,10 @@ def rows_read(shards, per_shard): def test_row_budget_keeps_a_floor_under_very_thin_shards(tmp_path): # Below the floor a file cannot witness the columns only it holds, which is the reason every file # is opened rather than a subset sampled. Overshooting the budget there is the right trade, and - # the profile says so: row_cap is 10, not the 1 the arithmetic asked for. + # the arithmetic share would be 1, so the floor holds and the budget is deliberately exceeded. for shard in range(10): _write_parquet(tmp_path / f"train-{shard:05d}-of-00010.parquet", [{"a": i} for i in range(50)]) result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=10) - assert {f.row_cap for s in result.partitions[0].splits for f in s.files} == {10} - assert result.sampling.rows_scanned == 100 # deliberately over the budget + assert result.sampling.rows_scanned == 100 # 10 files x the 10-row floor, over the budget of 10 From 377c2befbe5ff4a9c2b631a8acc0305f5517a5e6 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 6 Aug 2026 17:40:28 -0400 Subject: [PATCH 30/44] refactor(datasets): drop SamplingInfo.seed and FeatureSchema.fixed_length `seed` was never written. The producer passed `seed=None` with the comment "head sampling makes no random choices; a seed would be theatre", which was true and made the field dead on arrival. A test was still passing `seed=7` to the constructor, which pydantic silently ignored -- so nothing noticed either. `fixed_length` recorded the constant element count of a fixed-size list, e.g. an embedding's 768. Removed on request; noting the case against, since the measurement went the other way from my own earlier argument. It costs 19 characters, once per fixed-size-list column, and nothing at all for a dataset without one -- I had guessed it "rides the most-repeated node in the tree", but features are per column, not per file or per row. And it is the one thing in `features` that is not derivable from the rest of the profile: "list of float32" and "768-dimensional embedding" are different facts, and the second is exactly what a reader would otherwise fetch the data to learn. Restoring it is one field plus the two lines in schema.py that computed it. The schema tests that asserted it now assert what survives: a fixed-size list is still a list of its element type, indistinguishable from a variable one, so neither case can silently lose `items`. The contract is 67 fields across 15 models. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 14 ++------ .../tests/files/test_dataset_profile.py | 5 ++- .../nemo_datasets_plugin/profiler/pipeline.py | 1 - .../nemo_datasets_plugin/profiler/schema.py | 11 ++---- plugins/nemo-datasets/tests/test_schema.py | 34 +++++++------------ 5 files changed, 20 insertions(+), 45 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 38746fed98..d635eadc5a 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -228,13 +228,6 @@ class FeatureSchema(BaseModel): "as a suggestion to correct." ), ) - fixed_length: int | None = Field( - default=None, - description=( - "dtype == list: constant observed element count (e.g. an embedding vector's 768), None when " - "variable. Multi-dimensional shapes compose via nesting." - ), - ) fields: list[FeatureSchema] | None = Field(default=None, description="dtype == struct: named child fields.") items: FeatureSchema | None = Field(default=None, description="dtype in {list, messages}: element schema.") @@ -243,9 +236,9 @@ def _fields_and_items_are_exclusive(self) -> FeatureSchema: """A node is either a named-field container or has a single element schema, never both. Deliberately the only structural check here: it holds for *any* dtype, so it costs no - forward compatibility. Tying `fields` / `items` / `fixed_length` to specific dtype values - would instead reject a profile written by a newer profiler that added a container dtype, - which is exactly what the open vocabulary exists to prevent. + forward compatibility. Tying `fields` / `items` to specific dtype values would instead + reject a profile written by a newer profiler that added a container dtype, which is exactly + what the open vocabulary exists to prevent. """ if self.fields is not None and self.items is not None: raise ValueError(f"feature {self.name!r}: `fields` and `items` are mutually exclusive") @@ -478,7 +471,6 @@ class SamplingInfo(BaseModel): "witnesses, so a partition with very many files may exceed its budget." ), ) - seed: int | None = Field(default=None, description="RNG seed used for row selection, for reproducibility.") class DatasetProfile(BaseModel): diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index b35db8ed60..12207e9f5b 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -170,7 +170,6 @@ def _build_profile() -> DatasetProfile: files_read=2, files_present=2, row_budget=1024, - seed=7, ), partitions=[ PartitionProfile( @@ -322,8 +321,8 @@ def test_container_shape_is_not_pinned_to_known_dtypes(): profiler still loads on an older reader, which is what the open vocabulary buys.""" future_map = FeatureSchema(name="attrs", dtype="map", fields=[FeatureSchema(name="k", dtype="string")]) assert [field.name for field in future_map.fields or []] == ["k"] - future_tensor = FeatureSchema(name="embedding", dtype="tensor", fixed_length=768) - assert future_tensor.fixed_length == 768 + future_tensor = FeatureSchema(name="embedding", dtype="tensor", items=FeatureSchema(dtype="float32")) + assert future_tensor.items is not None and future_tensor.items.dtype == "float32" def test_unknown_fields_are_ignored_for_forward_compat(): diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 20455e1419..aede848f95 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -132,7 +132,6 @@ def profile( # than a bare count. Non-data files (a README, a LICENSE) are not data and are counted nowhere. files_present=len(data_entries) + sum(1 for e in file_errors if detect_format(e.path) is None), row_budget=row_budget, - seed=None, # head sampling makes no random choices; a seed would be theatre ) return DatasetProfile( created_at=created_at, diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py index 92d4c2c508..f5e5ccbb14 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py @@ -7,7 +7,7 @@ carries a declared schema, so it is converted directly; formats without one (jsonl) are inferred from the sampled rows by resolving each column's dtype. A list of ``{role, content}`` structs — or ShareGPT's ``{from, value}`` spelling of the same thing — is recognized as the ``messages`` dtype, -and a list whose elements are all the same length records that length as ``fixed_length``. +and a list of ``{role, content}`` structs is recognized as the ``messages`` dtype. """ from __future__ import annotations @@ -75,10 +75,7 @@ def _feature_from_arrow(name: str, arrow_type: pa.DataType) -> FeatureSchema: for i in range(arrow_type.num_fields) ] return FeatureSchema(name=name, dtype="struct", fields=fields) - if pa.types.is_fixed_size_list(arrow_type): - item = _feature_from_arrow("", arrow_type.value_type) - return FeatureSchema(name=name, dtype="list", items=item, fixed_length=arrow_type.list_size) - if pa.types.is_list(arrow_type) or pa.types.is_large_list(arrow_type): + if pa.types.is_list(arrow_type) or pa.types.is_large_list(arrow_type) or pa.types.is_fixed_size_list(arrow_type): item = _feature_from_arrow("", arrow_type.value_type) dtype = "messages" if _is_message_struct(item) else "list" return FeatureSchema(name=name, dtype=dtype, items=item) @@ -119,9 +116,7 @@ def _infer_feature(name: str, values: list[Any]) -> FeatureSchema: item = _infer_feature("", [element for value in present for element in value]) if _is_message_struct(item): return FeatureSchema(name=name, dtype="messages", items=item) - lengths = {len(value) for value in present} - fixed_length = lengths.pop() if len(lengths) == 1 else None - return FeatureSchema(name=name, dtype="list", items=item, fixed_length=fixed_length) + return FeatureSchema(name=name, dtype="list", items=item) return FeatureSchema(name=name, dtype=_scalar_dtype(present)) diff --git a/plugins/nemo-datasets/tests/test_schema.py b/plugins/nemo-datasets/tests/test_schema.py index e492ffcfe2..e8fd3364fd 100644 --- a/plugins/nemo-datasets/tests/test_schema.py +++ b/plugins/nemo-datasets/tests/test_schema.py @@ -25,19 +25,14 @@ def test_from_arrow_list_of_role_content_structs_is_messages(): assert [f.name for f in feature.items.fields] == ["role", "content"] -def test_from_arrow_fixed_size_list_records_length(): - schema = pa.schema([("embedding", pa.list_(pa.float32(), 768))]) - feature = derive_features([], schema)[0] - assert feature.dtype == "list" - assert feature.fixed_length == 768 - assert feature.items.dtype == "float32" - +def test_from_arrow_fixed_and_variable_lists_agree_on_shape(): + # A fixed-size list is still a list of its element type. The constant length itself is no longer + # recorded, so the two cases must be indistinguishable rather than one silently losing `items`. + fixed = derive_features([], pa.schema([("embedding", pa.list_(pa.float32(), 768))]))[0] + assert (fixed.dtype, fixed.items.dtype) == ("list", "float32") -def test_from_arrow_variable_list_has_no_fixed_length(): - feature = derive_features([], pa.schema([("tags", pa.list_(pa.string()))]))[0] - assert feature.dtype == "list" - assert feature.fixed_length is None - assert feature.items.dtype == "string" + variable = derive_features([], pa.schema([("tags", pa.list_(pa.string()))]))[0] + assert (variable.dtype, variable.items.dtype) == ("list", "string") # --- inferred from sampled rows (jsonl) ---------------------------------------------------------- @@ -70,17 +65,12 @@ def test_from_arrow_sharegpt_from_value_is_messages(): assert derive_features([], schema)[0].dtype == "messages" -def test_from_rows_constant_length_list_records_fixed_length(): - feature = derive_features([{"e": [0.1, 0.2, 0.3]}, {"e": [0.4, 0.5, 0.6]}])[0] - assert feature.dtype == "list" - assert feature.fixed_length == 3 - assert feature.items.dtype == "float64" - +def test_from_rows_lists_infer_their_element_type(): + constant = derive_features([{"e": [0.1, 0.2, 0.3]}, {"e": [0.4, 0.5, 0.6]}])[0] + assert (constant.dtype, constant.items.dtype) == ("list", "float64") -def test_from_rows_variable_length_list_has_no_fixed_length(): - feature = derive_features([{"e": [1, 2]}, {"e": [1, 2, 3]}])[0] - assert feature.dtype == "list" - assert feature.fixed_length is None + variable = derive_features([{"e": [1, 2]}, {"e": [1, 2, 3]}])[0] + assert (variable.dtype, variable.items.dtype) == ("list", "int64") def test_from_rows_nested_struct(): From db0df002fc2722d38592f323171bffad63816364 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 6 Aug 2026 23:22:58 -0400 Subject: [PATCH 31/44] feat(datasets): report how much a dataset weighs `size_bytes` came off FileRecord two commits ago on the reasoning that it "existed for the content digest, which was removed several commits ago and left [it] unread". The first half was true and the second was not: bytes also answer "will this fit?", which is the first question asked of an unfamiliar dataset and the one a row count cannot answer, since a row ranges from an integer score to a 32k-char reasoning trace. Dropping the per-file list was right; dropping the sum was collateral. Hugging Face publishes three of these -- num_bytes_original_files, num_bytes_parquet_files, num_bytes_memory -- at dataset, config and split granularity. It is the one thing their /info answers that ours could not answer at all. So SplitProfile.size_bytes, summed from the file listing. One integer per split, flat in shard count, which is the property the last commit was protecting. It is never None, and that is why it is a separate field rather than another nullable one alongside num_examples: size is read off the listing and a row count off the data, so the two go unknown independently. A shard that will not parse still weighs what it weighs. SamplingInfo.bytes_present covers what the splits cannot. A file in a format with no reader never reaches a partition, so a directory of .csv shards beside one .parquet would otherwise weigh in at the parquet alone -- the same "profiles as empty" failure file_errors exists to prevent, in the size column. It is redundant with the sum over splits exactly when nothing failed and load-bearing when something did, which is the property files_present already has and is kept for. Computing it made `unreadable_entries` an explicit list, which also stops files_present re-deriving a count it already had by running detect_format back over the accumulated errors. Not included: a decoded in-memory size. HF's num_bytes_memory is Arrow buffers (2.15x on-disk for OpenMathReasoning); ours would be Python objects, which the row-budget docstring guesses at 20x. Measuring that is a real pass over the data, not a number read off a listing, and is a separate question. The three golden fixtures carry real byte counts from HF's /size API, whose row counts they already matched exactly. The contract is 69 fields across 16 models. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 25 +++++++++++++ .../tests/files/test_dataset_profile.py | 37 ++++++++++++++----- .../nemo_datasets_plugin/profiler/pipeline.py | 22 ++++++++--- plugins/nemo-datasets/tests/test_pipeline.py | 24 ++++++++++++ 4 files changed, 94 insertions(+), 14 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index d635eadc5a..34f9e5cd38 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -344,6 +344,19 @@ class SplitProfile(BaseModel): "grows without bound and informs no decision." ), ) + size_bytes: int = Field( + default=0, + description=( + "On-disk bytes of this split's files, summed. Answers whether the data fits wherever the " + "reader means to put it — the first question asked of an unfamiliar dataset, and one a row " + "count cannot answer, since a row ranges from an integer score to a reasoning trace. " + "Unlike `num_examples` this is never None: it comes from the file listing rather than from " + "reading, so a file that failed mid-read still contributes its size. Bytes as stored — " + "compressed, and several times this once decoded into memory. Covers only files a " + "partition grouped; a format with no reader never reaches a split, so weigh the whole " + "fileset with `SamplingInfo.bytes_present`." + ), + ) num_examples: int | None = Field( default=None, description=( @@ -460,6 +473,18 @@ class SamplingInfo(BaseModel): "these two to match until scale forces file-level sampling." ), ) + bytes_present: int = Field( + default=0, + description=( + "On-disk bytes of every data file the fileset holds, whether or not this run could read it " + "— the size of the dataset as it sits, independent of how much was profiled. Redundant " + "with the sum over `SplitProfile.size_bytes` exactly when nothing failed, and load-bearing " + "when something did: a file in a format with no reader never reaches a partition, so a " + "directory of .csv shards beside one .parquet would otherwise weigh in at the parquet " + "alone. Same reason `files_present` is kept alongside the per-split counts — a denominator " + "stops being derivable the moment coverage is partial, which is the only time it is read." + ), + ) row_budget: int | None = Field( default=None, description=( diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index 12207e9f5b..700fc58053 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -36,14 +36,14 @@ created_at: 2026-07-08T22:05:12Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 2112, rows_present: 3201061, - files_read: 33, files_present: 33, row_budget: 4096} + files_read: 33, files_present: 33, bytes_present: 31821490182, row_budget: 4096} partitions: - name: "" file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 3200861, num_files: 32} - - {name: test, canonical: test, num_examples: 200, num_files: 1} + - {name: train, canonical: train, num_examples: 3200861, num_files: 32, size_bytes: 31819412254} + - {name: test, canonical: test, num_examples: 200, num_files: 1, size_bytes: 2077928} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -75,14 +75,14 @@ created_at: 2026-07-08T22:41:37Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 1024, rows_present: 46189, - files_read: 2, files_present: 2, row_budget: 1024} + files_read: 2, files_present: 2, bytes_present: 27055195, row_budget: 1024} partitions: - name: "" file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 43835, num_files: 1} - - {name: test, canonical: test, num_examples: 2354, num_files: 1} + - {name: train, canonical: train, num_examples: 43835, num_files: 1, size_bytes: 25670988} + - {name: test, canonical: test, num_examples: 2354, num_files: 1, size_bytes: 1384207} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -114,14 +114,14 @@ created_at: 2026-07-09T10:12:45Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 1024, rows_present: 21362, - files_read: 2, files_present: 2, row_budget: 1024} + files_read: 2, files_present: 2, bytes_present: 19459677, row_budget: 1024} partitions: - name: "" file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 20324, num_files: 1} - - {name: validation, canonical: validation, num_examples: 1038, num_files: 1} + - {name: train, canonical: train, num_examples: 20324, num_files: 1, size_bytes: 18495985} + - {name: validation, canonical: validation, num_examples: 1038, num_files: 1, size_bytes: 963692} features: - {name: prompt, dtype: string, semantic_role: prompt, semantic_role_source: detected} - {name: response, dtype: string, semantic_role: completion, semantic_role_source: detected} @@ -250,6 +250,25 @@ def test_openmathreasoning_locks_contract_shape(): assert part.stats["completion"].messages.roles_seen == ["assistant"] +@pytest.mark.parametrize("name", sorted(FIXTURES)) +def test_split_sizes_account_for_the_whole_fileset(name): + """On a clean profile the splits weigh the whole fileset, so `bytes_present` is the same number + reached without going through partitions. That redundancy is the point: it is what lets the + figure survive a file no partition could group.""" + profile = DatasetProfile.model_validate(yaml.safe_load(FIXTURES[name])) + assert not profile.file_errors + from_splits = sum(split.size_bytes for part in profile.partitions for split in part.splits) + assert from_splits == profile.sampling.bytes_present + + +def test_a_split_weighs_something_even_when_its_row_count_does_not(): + """Size is read off the file listing and a row count off the data, so they go unknown + independently — `num_examples` is None-able and `size_bytes` is not.""" + split = SplitProfile(name="train", num_files=3, size_bytes=4096) + assert split.num_examples is None + assert split.size_bytes == 4096 + + def test_helpsteer2_flat_schema_and_no_verifiability(): profile = DatasetProfile.model_validate(yaml.safe_load(HELPSTEER2)) part = profile.partitions[0] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index aede848f95..9e8501ca51 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -94,15 +94,20 @@ def profile( # Files that plainly hold records but have no reader yet. They are not profiled, but they must be # reported: silently dropping them let a directory of .csv shards profile as an exhaustively # scanned, empty dataset — indistinguishable from a dataset that really is empty. They get real - # FileRecords like any other file the profiler could not read, just at the envelope, since no - # partition ever grouped them. + # FileErrors like any other file the profiler could not read, just at the envelope, since no + # partition ever grouped them. Kept as entries, not just paths, because their bytes still count + # toward the size of the fileset even though no partition will ever weigh them. + unreadable_entries = [ + entry + for entry in sorted(all_entries, key=lambda entry: entry.path) + if detect_format(entry.path) is None and is_unsupported_data(entry.path) + ] file_errors = [ FileError( path=entry.path, error=f"no reader for '{PurePosixPath(entry.path).suffix.lower()}' files", ) - for entry in sorted(all_entries, key=lambda entry: entry.path) - if detect_format(entry.path) is None and is_unsupported_data(entry.path) + for entry in unreadable_entries ] partitions: list[PartitionProfile] = [] @@ -130,7 +135,11 @@ def profile( files_read=files_read, # files actually opened and read, not files merely listed # Every data file, readable or not: the denominator that makes `files_read` a fraction rather # than a bare count. Non-data files (a README, a LICENSE) are not data and are counted nowhere. - files_present=len(data_entries) + sum(1 for e in file_errors if detect_format(e.path) is None), + files_present=len(data_entries) + len(unreadable_entries), + # Weighed over the same set, so a fileset the profiler could not read still reports its size. + # Summing the splits would miss the unreadable files, which never reach a partition. + bytes_present=sum(entry.size_bytes for entry in data_entries) + + sum(entry.size_bytes for entry in unreadable_entries), row_budget=row_budget, ) return DatasetProfile( @@ -321,6 +330,9 @@ def _profile_partition( name=split.name, canonical=split.canonical, num_files=len(split.entries), + # From the listing, not from reading, so a file that failed mid-read still weighs + # what it weighs — unlike `num_examples`, this never goes unknown. + size_bytes=sum(entry.size_bytes for entry in split.entries), num_examples=split_examples if split_counts_known else None, ) ) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 2597f86d9a..8addadbd6e 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -346,6 +346,11 @@ def test_profile_reports_unsupported_data_files(tmp_path): assert result.sampling.rows_present is None # not 0: "empty" would be a lie assert result.sampling.files_read == 0 assert result.sampling.files_present == 2 # both are data; neither could be read + # ...and they still weigh what they weigh. This is the case `bytes_present` exists for: no + # partition grouped these files, so summing the splits reports zero -- the same lie as "empty". + on_disk = (tmp_path / "train.csv").stat().st_size + (tmp_path / "test.arrow").stat().st_size + assert result.sampling.bytes_present == on_disk + assert sum(s.size_bytes for p in result.partitions for s in p.splits) == 0 # Typed records now, each saying why -- not bare paths tucked into a free-form dict. assert [e.path for e in result.file_errors] == ["test.arrow", "train.csv"] assert all("no reader" in e.error for e in result.file_errors) @@ -362,6 +367,25 @@ def test_profile_ignores_non_data_files_without_penalty(tmp_path): assert result.partitions[0].stats_complete is True assert result.file_errors == [] assert result.sampling.files_present == 1 # the README and LICENSE are not data, counted nowhere + # Nor does their weight land on the dataset: a card is not part of what has to be moved. + assert result.sampling.bytes_present == (tmp_path / "train-00000-of-00001.parquet").stat().st_size + + +def test_split_size_survives_a_shard_it_could_not_read(tmp_path): + # Size comes from the listing and a row count from reading, so they go unknown independently. + # A shard that fails to parse still weighs what it weighs, where the split's row count cannot. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"a": 1}, {"a": 2}]) + (tmp_path / "train-00001-of-00002.parquet").write_bytes(b"not parquet") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + on_disk = sum(p.stat().st_size for p in tmp_path.glob("*.parquet")) + split = result.partitions[0].splits[0] + assert split.num_files == 2 + assert split.size_bytes == on_disk # both shards, including the one that would not open + assert split.num_examples is None # the broken shard's rows are unknowable... + assert split.size_bytes > 0 # ...but its bytes are not + assert result.sampling.bytes_present == on_disk def test_profile_records_a_partial_jsonl_read(tmp_path): From fb2ca48f126e321eeada7ba2516ee59488551b5c Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 11:45:26 -0400 Subject: [PATCH 32/44] docs(datasets): correct the pipeline module docstring Three claims in it were false, one of them describing the exact bug the branch fixed. "partitions, splits, FileRecords, content digest, sampling metadata" -- FileRecord became FileError and the content digest was removed, so the docstring named two types that no longer exist and omitted the split-level counts and sizes that replaced them. "each is read up to `row_cap` rows, so peak memory tracks the file count rather than the dataset size" -- that is the per-file cap, which was replaced by a per-partition budget precisely because peak memory tracking the file count is the failure mode: resharding a dataset multiplied heap while describing the same data. The docstring was still recommending the behaviour as the design. "Pass `row_cap=None`" -- the parameter is `row_budget`. `row_cap` survives as the per-file share `_per_file_cap` computes and the readers take, so the name is not stale everywhere, only here where it stood in for the public knob. Now states the budget's actual axis, points at DEFAULT_ROW_BUDGET for the measurement rather than repeating it, and records MIN_ROWS_PER_FILE as what makes the budget a target rather than a ceiling. Scoped to this docstring: the other `exhaustive` mentions in the profiler are the English adjective, not the deleted SamplingInfo field. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/pipeline.py | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 9e8501ca51..2c009479ab 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -4,15 +4,22 @@ """The top-level profiling pipeline. ``profile(source)`` lists the files behind a :class:`FileSource`, groups them into partitions and -splits, reads them, and assembles a ``DatasetProfile``. This stage produces the structural envelope -— partitions, splits, FileRecords, content digest, sampling metadata — the derived row schema -(``features``), per-column ``stats``, and the full ``classification`` (roles, format, prompt form, -dataset type, and verifiability). - -Every file is opened — sampling a subset of files would hide columns that appear only in later shards -— but each is read up to ``row_cap`` rows, so peak memory tracks the file count rather than the -dataset size. Files smaller than the cap are read to the end and keep their exact counts, so capping -costs nothing on a small dataset. Pass ``row_cap=None`` for a genuinely exhaustive scan. +splits, reads them, and assembles a ``DatasetProfile``. It produces the structural envelope — +partitions, splits with their counts and sizes, the files it could not use, and the sampling figures +— along with the derived row schema (``features``), per-column ``stats``, and the full +``classification`` (roles, format, prompt form, dataset type, and verifiability). + +Every file is opened — sampling a *subset of files* would hide columns that appear only in later +shards — but a partition's ``row_budget`` is divided across its files, so peak memory tracks the +budget rather than the shard count. Capping each file instead put the knob on the wrong axis: +resharding the same data then multiplied the rows held in memory without describing any more of it. +See :data:`DEFAULT_ROW_BUDGET`. + +The budget is a target rather than a ceiling. :data:`MIN_ROWS_PER_FILE` is the floor every file is +read to however thin its share gets, since one sampled below it cannot contribute the columns it +alone witnesses. Files smaller than their share are read to the end and keep exact row counts, so a +budgeted profile of a small dataset is still complete. Pass ``row_budget=None`` for a genuinely +exhaustive scan. """ from __future__ import annotations From be6130fac036b40885489c7d1a8525cb10d7df76 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 14:26:41 -0400 Subject: [PATCH 33/44] feat(datasets): infer a glob that selects each split's files `SplitProfile.data_files` -- one pattern per split, relative to the fileset root: helpsteer2/train*.parquet helpsteer2/validation*.parquet Gives the files back their addressability without giving back the per-file manifest that was removed for scaling. A consumer can hand a reader the files of one split without listing the fileset and re-deriving which shards belong where, and it stays one string whether the split has 1 shard or 10,000. Inferred as the inverse of split resolution: splits are read *off* the paths, so the pattern is rebuilt from the same evidence -- the directory the files share and the split name their filenames start with. The two common layouts fall out naturally and want opposite candidates: data/train-00000-of-00143.parquet -> data/train*.parquet default/train/0000.parquet -> default/train/*.parquet Named `data_files` for HF card front-matter's `configs[].data_files`, which is the declared form of this same claim. SplitProfile already documents card-declared splits as resolution tier 1, so when cards are parsed the declared value replaces this inference rather than sitting beside it in a second vocabulary. (Not `files`, which was the deleted per-file list -- a reader seeing that name back would expect a list, and this is one string.) Never approximate. Every candidate is matched back against every file the source listed -- not just the split's, not just the partition's -- and only one that reproduces the split exactly is emitted. A glob is an instruction to go read files, so a near miss is not a rougher answer: `data/*` over a directory holding a README hands a card to a trainer as though it were a shard. None is a first-class result and means "not expressible as one pattern", which a consumer can handle. Verification is also what resolves sibling splits. `train` beside `train_prefs` makes `train*` match both, so the simple form loses and the narrower `train-*` is reached -- candidates run simplest-first precisely so the tidy pattern is what appears unless something forces otherwise. Only `*`, never `**`. One reading of `*` -- any run of characters except `/` -- is shared by shell globs, Python's glob, fsspec and HF, so a pattern means the same thing wherever it is pasted. `**` does not have that property, so a split whose shards span subdirectories reports None instead. A test resolves every emitted pattern through pathlib's own glob rather than the matcher here, since the claim is about what a *consumer* will select. The contract is 70 fields across 16 models. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 20 ++++ .../tests/files/test_dataset_profile.py | 37 +++++-- .../nemo_datasets_plugin/profiler/pipeline.py | 16 ++- .../nemo_datasets_plugin/profiler/splits.py | 72 ++++++++++++++ plugins/nemo-datasets/tests/test_pipeline.py | 97 ++++++++++++++++++- 5 files changed, 233 insertions(+), 9 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 34f9e5cd38..30f6458ec8 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -335,6 +335,26 @@ class SplitProfile(BaseModel): "train, with the variant's intent kept in `name`." ), ) + data_files: str | None = Field( + default=None, + description=( + "A glob selecting exactly this split's files, relative to the fileset root: \"helpsteer2/" + 'train*.parquet". Gives the files back their addressability without giving back the ' + "per-file manifest — one pattern per split, whatever the shard count — so a consumer can " + "hand a reader the files of one split without listing the fileset and re-deriving which " + "shards belong where. Named for HF card front-matter's `configs[].data_files`, which is " + "the declared form of this same claim and, once cards are parsed, the thing that will " + "replace this inference rather than sit beside it in a second vocabulary.\n\n" + "`*` spans any run of characters except `/` — the one reading shared by shell globs, " + "Python's glob, fsspec and HF — so the pattern means the same thing wherever it is pasted. " + "`**` is never emitted, because its meaning is not shared.\n\n" + "None when no single pattern selects these files and nothing else (shards spread across " + "subdirectories, say). Never approximate: a pattern is emitted only after being matched " + "back against every file in the fileset and found to select this split exactly. A glob is " + "an instruction to go read files, so a near miss is not a rougher answer — it silently " + "pulls a README, or a neighbouring split's shards, into a training set." + ), + ) num_files: int = Field( default=0, description=( diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index 700fc58053..9bf0efcf59 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -42,8 +42,10 @@ file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 3200861, num_files: 32, size_bytes: 31819412254} - - {name: test, canonical: test, num_examples: 200, num_files: 1, size_bytes: 2077928} + - {name: train, canonical: train, num_examples: 3200861, num_files: 32, + size_bytes: 31819412254, data_files: 'train*.parquet'} + - {name: test, canonical: test, num_examples: 200, num_files: 1, + size_bytes: 2077928, data_files: 'test*.parquet'} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -81,8 +83,10 @@ file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 43835, num_files: 1, size_bytes: 25670988} - - {name: test, canonical: test, num_examples: 2354, num_files: 1, size_bytes: 1384207} + - {name: train, canonical: train, num_examples: 43835, num_files: 1, + size_bytes: 25670988, data_files: 'train*.parquet'} + - {name: test, canonical: test, num_examples: 2354, num_files: 1, + size_bytes: 1384207, data_files: 'test*.parquet'} features: - {name: prompt, dtype: messages, semantic_role: prompt, semantic_role_source: detected, items: {dtype: struct, fields: [{name: role, dtype: string}, {name: content, dtype: string}]}} @@ -120,8 +124,10 @@ file_formats: [parquet] stats_complete: false splits: - - {name: train, canonical: train, num_examples: 20324, num_files: 1, size_bytes: 18495985} - - {name: validation, canonical: validation, num_examples: 1038, num_files: 1, size_bytes: 963692} + - {name: train, canonical: train, num_examples: 20324, num_files: 1, + size_bytes: 18495985, data_files: 'train*.parquet'} + - {name: validation, canonical: validation, num_examples: 1038, num_files: 1, + size_bytes: 963692, data_files: 'validation*.parquet'} features: - {name: prompt, dtype: string, semantic_role: prompt, semantic_role_source: detected} - {name: response, dtype: string, semantic_role: completion, semantic_role_source: detected} @@ -261,6 +267,25 @@ def test_split_sizes_account_for_the_whole_fileset(name): assert from_splits == profile.sampling.bytes_present +@pytest.mark.parametrize("name", sorted(FIXTURES)) +def test_split_globs_are_one_pattern_each_and_never_cross_a_directory(name): + """`data_files` is a single pattern, not a manifest, so it cannot reintroduce the per-file growth + the split-level counts exist to avoid. `**` is never emitted, because its meaning is not shared + across glob implementations.""" + profile = DatasetProfile.model_validate(yaml.safe_load(FIXTURES[name])) + for part in profile.partitions: + for split in part.splits: + assert isinstance(split.data_files, str) + assert "**" not in split.data_files + + +def test_a_split_with_no_expressible_pattern_says_so(): + """None is a first-class answer: shards spread across subdirectories need `**` to cover, and a + pattern that resolves differently in the reader than in the profiler is worse than none.""" + split = SplitProfile(name="train", num_files=2) + assert split.data_files is None + + def test_a_split_weighs_something_even_when_its_row_count_does_not(): """Size is read off the file listing and a row count off the data, so they go unknown independently — `num_examples` is None-able and `size_bytes` is not.""" diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 2c009479ab..5676663374 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -34,7 +34,7 @@ from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data from nemo_datasets_plugin.profiler.schema import derive_features -from nemo_datasets_plugin.profiler.splits import resolve_splits +from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats, quote_enumerations from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, @@ -123,8 +123,15 @@ def profile( # None once any file's row count is unknown: the fileset's total is then unknowable, not zero. rows_present: int | None = 0 + # Every path the source listed, data or not. A split's glob is verified against this rather than + # against the partition's own files, so a pattern can never be emitted that would also pull in a + # README sitting beside the shards. + all_paths = [entry.path for entry in all_entries] + for name, partition_entries in group_partitions(data_entries): - outcome = _profile_partition(source, name, partition_entries, row_budget, column_roles or {}) + outcome = _profile_partition( + source, name, partition_entries, row_budget, column_roles or {}, all_paths=all_paths + ) partitions.append(outcome.partition) rows_scanned += outcome.rows_scanned files_read += outcome.files_read @@ -265,6 +272,8 @@ def _profile_partition( entries: list[FileEntry], row_budget: int | None, column_roles: dict[str, str], + *, + all_paths: list[str], ) -> _PartitionOutcome: """Profile one partition — the files of one source directory, whatever formats they are in. @@ -336,6 +345,9 @@ def _profile_partition( SplitProfile( name=split.name, canonical=split.canonical, + # Inferred from the same paths the split itself was read off, then verified against + # the whole listing; None when one pattern cannot express the split exactly. + data_files=infer_data_files(split.name, split.entries, all_paths), num_files=len(split.entries), # From the listing, not from reading, so a file that failed mid-read still weighs # what it weighs — unlike `num_examples`, this never goes unknown. diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py index 6b5c1afc9a..dfb79c45ba 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/splits.py @@ -75,6 +75,78 @@ def _split_name(path: str) -> str: return _SHARD_SUFFIX.sub("", parts[-1].split(".")[0]) +def _glob_matches(pattern: str, path: str) -> bool: + """Match ``path`` against ``pattern`` where ``*`` spans any run of characters except ``/``. + + Deliberately the narrowest dialect rather than the most expressive one. This single reading of + ``*`` is shared by shell globs, Python's :mod:`glob`, fsspec and HF ``data_files``, so a pattern + emitted here means the same thing wherever a consumer pastes it. ``**`` is not produced at all: + its semantics differ between those tools, and a pattern that silently selects a different set of + files in the reader than it did in the profiler is worse than no pattern. + """ + return re.fullmatch("[^/]*".join(re.escape(part) for part in pattern.split("*")), path) is not None + + +def infer_data_files(split_name: str, entries: list[FileEntry], all_paths: list[str]) -> str | None: + """A glob selecting exactly ``entries`` out of ``all_paths``, or None if no single one does. + + This is the inverse of the split resolution above: splits are read *off* the paths, so the + pattern is rebuilt from the same evidence — the directory the files share, and the split name + their filenames start with. It restores addressability of a split's files without restoring the + per-file manifest that was removed for scaling: one pattern per split, whatever the shard count. + + Candidates run most specific first, because the general ones over-match. In ``helpsteer2/`` a + ``train`` split wants ``helpsteer2/train*.parquet``; ``helpsteer2/*.parquet`` would swallow the + validation shards too, and only the ordering distinguishes them. The ``data/train/0000.parquet`` + layout is the other way round — no filename starts with "train", so the name-anchored candidates + are skipped and the directory-wide one is exactly right. + + Every candidate is matched back against *every* file the source listed, not just this split's or + even just this partition's, and the first that reproduces the split exactly wins. Nothing is + emitted on a near miss. A glob is an instruction to go read files, so an approximate one is not + a smaller version of the right answer -- it quietly pulls a README, or another split's shards, + into a training set. None says "these files are not expressible as one pattern", which is a + thing a consumer can handle; a wrong pattern is not. + """ + paths = [entry.path for entry in entries] + directories = {PurePosixPath(path).parent.as_posix() for path in paths} + if len(directories) != 1: + # Shards spread across subdirectories. Covering them needs `**`, whose meaning is not shared + # across glob implementations, so this reports no pattern rather than an ambiguous one. + return None + directory = directories.pop() + prefix = "" if directory == "." else f"{directory}/" + names = [PurePosixPath(path).name for path in paths] + suffixes = {PurePosixPath(path).suffix for path in paths} + # A partition may hold more than one format; only a single shared suffix can go in the pattern. + suffix = suffixes.pop() if len(suffixes) == 1 else None + + # Stems to anchor on: the bare split name first, then the same name keeping the separator that + # follows it. Plain `train*` is what a person would write and is tried first for that reason; + # the separator variants exist only for the sibling collision, where `train` sits beside + # `train_prefs` in one directory and `train*` swallows both. Verification is what demotes the + # simple form there, so the narrower `train-*` is reached only when it is actually needed. Each + # stem is offered only when every filename in the split carries it, so one can never exclude a + # file it ought to match. + stems = [split_name] + [f"{split_name}{sep}" for sep in ("-", ".", "_")] if split_name else [] + candidates: list[str] = [] + for stem in stems: + if not all(name.startswith(stem) for name in names): + continue + if suffix: + candidates.append(f"{prefix}{stem}*{suffix}") + candidates.append(f"{prefix}{stem}*") + if suffix: + candidates.append(f"{prefix}*{suffix}") + candidates.append(f"{prefix}*") + + target = set(paths) + for candidate in candidates: + if {path for path in all_paths if _glob_matches(candidate, path)} == target: + return candidate + return None + + def resolve_splits(entries: list[FileEntry]) -> list[ResolvedSplit]: """Group files into splits by path inference. diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 8addadbd6e..692479c84d 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -5,13 +5,15 @@ import json from datetime import datetime, timezone +from pathlib import Path import pyarrow as pa import pyarrow.parquet as pq +import pytest from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.pipeline import _measure, profile -from nemo_datasets_plugin.profiler.splits import resolve_splits +from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits from nemo_platform_plugin.files.dataset_profile import DatasetProfile FIXED_TIME = datetime(2026, 7, 13, 12, 0, 0, tzinfo=timezone.utc) @@ -63,6 +65,99 @@ def test_resolve_splits_falls_back_to_single_default(): assert len(splits[0].entries) == 2 +# --- data_files glob inference -------------------------------------------------------------------- + + +def _globs(*paths): + """Infer a glob per split over ``paths``, verifying against the full listing (README included).""" + data = [e for e in _entries(*paths) if e.path.endswith((".parquet", ".jsonl"))] + return {s.name: infer_data_files(s.name, s.entries, list(paths)) for s in resolve_splits(data)} + + +@pytest.mark.parametrize( + "label,paths,expected", + [ + ( + "shards in one directory", + ( + "data/train-00000-of-00002.parquet", + "data/train-00001-of-00002.parquet", + "data/test-00000-of-00001.parquet", + ), + {"train": "data/train*.parquet", "test": "data/test*.parquet"}, + ), + ( + "a directory per split", + ("default/train/0000.parquet", "default/train/0001.parquet", "default/test/0000.parquet"), + {"train": "default/train/*.parquet", "test": "default/test/*.parquet"}, + ), + ( + "files at the fileset root", + ("train.jsonl", "validation.jsonl"), + {"train": "train*.jsonl", "validation": "validation*.jsonl"}, + ), + ( + "no split detected: the glob covers the partition", + ("shard-00000.parquet", "shard-00001.parquet"), + {"default": "*.parquet"}, + ), + ( + "mixed formats drop the suffix rather than losing a file", + ("train-00000-of-00002.parquet", "train-00001-of-00002.jsonl"), + {"train": "train*"}, + ), + ], +) +def test_data_files_glob_per_layout(label, paths, expected): + assert _globs(*paths) == expected, label + + +def test_data_files_glob_excludes_a_readme_beside_the_shards(): + # `data/*` would sweep the card into the split. The suffix-qualified candidate is what survives + # verification, and verification runs against every listed file, not just the data ones. + assert _globs("data/train-00000-of-00001.parquet", "data/README.md") == {"train": "data/train*.parquet"} + + +def test_data_files_glob_keeps_a_separator_to_beat_a_sibling_split(): + # `train*` would also match train_prefs, so the simple form loses verification and the narrower + # `train-*` is reached. Both splits still get a pattern; neither over-matches the other. + assert _globs( + "train-00000-of-00002.parquet", "train-00001-of-00002.parquet", "train_prefs-00000-of-00001.parquet" + ) == {"train": "train-*.parquet", "train_prefs": "train_prefs*.parquet"} + + +def test_data_files_glob_refuses_rather_than_sweep_in_a_non_data_file(): + # Mixed suffixes leave no suffix to qualify with, and an unsplit-named set leaves no stem to + # anchor on, so the only candidate left is `data/*` -- which would hand a reader the README as + # if it were a shard. Verification is the whole of what stops that, and None is the answer. + assert _globs("data/shard-00000.parquet", "data/shard-00001.jsonl", "data/README.md") == {"default": None} + + +def test_data_files_glob_is_none_when_shards_span_subdirectories(): + # Expressing this needs `**`, whose meaning differs between glob implementations. None is the + # honest answer; a pattern that selects a different set in the reader than here would not be. + assert _globs("train/part-a/0000.parquet", "train/part-b/0000.parquet") == {"train": None} + + +def test_data_files_glob_means_the_same_thing_to_pythons_own_glob(tmp_path): + """The dialect claim, checked against an independent implementation rather than our matcher. + + A pattern is only worth emitting if a consumer resolves it to the files we said it selects. + """ + for rel in ( + "data/train-00000-of-00002.parquet", + "data/train-00001-of-00002.parquet", + "data/test-00000-of-00001.parquet", + ): + _write_parquet(tmp_path / rel, [{"a": 1}]) + (tmp_path / "data" / "README.md").write_text("card") + + for split in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0].splits: + resolved = sorted(p.relative_to(tmp_path).as_posix() for p in Path(tmp_path).glob(split.data_files)) + assert len(resolved) == split.num_files, f"{split.name}: {split.data_files} -> {resolved}" + assert all(name.startswith(f"data/{split.name}") for name in resolved) + + # --- partition grouping -------------------------------------------------------------------------- From 37f3ab23b3d80da3c1fbcc79ad0e0becb93b6dd1 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 16:48:30 -0400 Subject: [PATCH 34/44] feat(datasets): count distinct values only while a column is a vocabulary Phase 1 of the streaming spec. `CategoricalStats` is now emitted only when the column really is a bounded controlled vocabulary; above the bounds it is absent, and the absence is the claim. Counting distinct values exactly means retaining them, and on a free-text column the set of distinct values is the column. What it bought was "response: 9,954 distinct in 10,000 rows" -- which says free text, which `semantic_role` and the length quantiles already say for nothing. Nothing read it either: `distinct_count` has exactly two consumers, a `<= 2` test that confirms a binary preference label and the `<= 32` gate on quoting. For a string column the first can never fire, so every value above 32 was stored for a reader and read by no code. Three bounds, not one: >1024 distinct, any single value over 256 chars, or 64 KB retained. The middle one does the real work, and is the same kind of test as the role gate on quoting -- it asks what the column *is* rather than how many values it holds. A vocabulary member is short by nature, so one long value settles it on sight rather than after a thousand. `distinct_count` therefore stays a plain exact int. No None, no `saturated` flag, no estimate: the field exists only in the case where it is cheap and true. On the measured benefit, correcting my own spec. I claimed this removes ~73% of peak memory, from a deep-sizeof of the distinct sets. That double-counted: Python strings are shared, so today the set stores pointers into rows that are resident anyway, and the marginal cost is the hash table -- 2.6 MB beside 61.4 MB of rows, about 7% of peak. The bound is a precondition for streaming, not a win today. Once a batch is folded and discarded, a distinct set is the sole owner of every value it kept: the same two columns cost 46.8 MB, against 0.163 MB for every other accumulator combined. Unbounded cardinality is the single term that would drag the streaming fold back to O(rows). Landing it first because the contract change stands on its own merits and because Phase 3 cannot hold its invariant without it. Verified on the two-partition demo: the five HelpSteer2 rating columns keep their vocabularies (5 values each), prompt and response lose theirs, messages columns are unaffected -- they never had one, lists being unhashable. The contract is 70 fields across 16 models. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 28 +++++++--- .../nemo_datasets_plugin/profiler/stats.py | 55 +++++++++++++++++-- plugins/nemo-datasets/tests/test_stats.py | 48 +++++++++++++--- 3 files changed, 111 insertions(+), 20 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 30f6458ec8..9170ee5cff 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -246,18 +246,29 @@ def _fields_and_items_are_exclusive(self) -> FeatureSchema: class CategoricalStats(BaseModel): - """Cardinality signals for string / int columns. + """The vocabulary of a column that has one. - ``distinct_count`` is a count, not row content, and is always safe to store. The values - themselves ARE row content, so they appear only for a column whose detected role makes it a - controlled vocabulary — the assert-only-what-was-proven rule applied to the one place the + Present only when the column really is a bounded controlled vocabulary. Absent otherwise, and + the absence *is* the claim: this column is not a vocabulary. + + It used to be a general cardinality count on every string and numeric column. Counting distinct + values exactly means *retaining* them, and for a column of prompts the set of distinct values is + the column. What that bought was a reading of "9,954 distinct in 10,000 rows", which says free + text, which ``semantic_role`` and the length quantiles already said for free. Nothing read it + either: the only consumers of the number are a ``<= 2`` test that confirms a binary label and + the ``<= 32`` gate on ``values`` below. + + The values themselves ARE row content, so they appear only for a column whose detected role makes + it a controlled vocabulary — the assert-only-what-was-proven rule applied to the one place the profiler would otherwise leak the data it is describing. """ distinct_count: int = Field( description=( - "Distinct values among scanned rows: ~=rows_scanned -> id-like; a small bounded set " - "corroborates score / category roles." + "How many distinct values the vocabulary holds. Exact, with no cap to have silently hit: " + "this model is built only for a column that stayed inside the vocabulary bounds all the " + "way through, so there is nothing to caveat. A small bounded set corroborates score / " + "category roles, and `<= 2` is what confirms a binary preference label." ), ) values: list[str] | None = Field( @@ -287,7 +298,10 @@ class ColumnStats(BaseModel): text: TextStats | None = Field(default=None, description="dtype == string") numeric: NumericStats | None = Field(default=None, description="dtype in {int*, uint*, float*}") messages: MessageStats | None = Field(default=None, description="dtype == messages (list of {role, content})") - categorical: CategoricalStats | None = Field(default=None, description="low observed cardinality only") + categorical: CategoricalStats | None = Field( + default=None, + description="Present only when the column is a bounded controlled vocabulary; absence means it is not one.", + ) quality: TextQuality | None = Field(default=None, description="dtype == string: corruption signals") diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 9a3acaf03f..e44f763479 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -36,6 +36,19 @@ # A quotable enumeration holds at most this many distinct values. _MAX_ENUM_VALUES = 32 +# Where a column stops being a plausible controlled vocabulary, and so stops being worth counting. +# Three bounds because a count alone bounds cardinality but not bytes -- 1024 reasoning traces is +# 32 MB. `_MAX_VOCABULARY_VALUE_CHARS` is the one that matters most: it is a claim about what the +# column *is* rather than how big it is, so it settles a free-text column on the first value instead +# of after a thousand. Sized well above real vocabularies -- a `source` column spanning 500 datasets, +# a 200-class label set -- and far below anything that costs memory. +_MAX_VOCABULARY_VALUES = 1024 +_MAX_VOCABULARY_VALUE_CHARS = 256 +_MAX_VOCABULARY_BYTES = 64 * 1024 +# What a non-string distinct value is charged against the byte bound. Ints and bools are small and +# uniform, so an exact size buys nothing the count bound does not already give. +_NON_STRING_VALUE_BYTES = 8 + # Roles that are controlled vocabularies by construction, and so are safe to quote at any dataset # size. Everything else -- prompts, completions, chosen/rejected, context, chat -- is free text no # matter how few distinct values a small sample happens to show, and unroled columns are unknown, @@ -150,12 +163,42 @@ def at(percentile: int) -> int: def _cardinality(present: list[Any]) -> CategoricalStats | None: - """The count only. The values themselves are row data and are gated on role, not cardinality, - so :func:`quote_enumerations` adds them once classification has assigned one.""" - try: - distinct = set(present) - except TypeError: - return None # unhashable values (dicts / lists) have no cardinality signal + """The distinct-value count, for as long as the column looks like a controlled vocabulary. + + Returns None the moment it stops looking like one, discarding what it had accumulated. Counting + distinct values exactly means *retaining* them, so on a free-text column this set grows to hold + the column, to report that `response` had 9,954 distinct values in 10,000 rows -- which is "this + is free text", which the role marker and the length quantiles already say for nothing. + + Today that costs little, because the rows are held anyway and the set stores pointers into them: + measured at 2.6 MB beside 61.4 MB of resident rows. The bound is here for the pipeline this is + becoming. Once a batch is folded and discarded, a distinct set is the *sole owner* of every value + it kept, and the same two columns cost 46.8 MB against 0.163 MB for every other accumulator + combined. Unbounded, this is the one thing that would make a streaming fold O(rows) again. + + Three bounds rather than one. A count alone bounds cardinality but not bytes, and 1024 reasoning + traces is 32 MB. The middle bound is the one that does the real work: it asks what the column + *is* rather than how many values it holds, in the same way the role gate on quoting does. A + vocabulary member is short by nature, so a single 32 KB value settles the question on sight, + which is why free-text columns exit here almost immediately instead of after 1024 values. + + The values themselves are still never returned here. They are row content, gated on role rather + than on size, and :func:`quote_enumerations` adds them once classification has assigned one. + """ + distinct: set[Any] = set() + retained_bytes = 0 + for value in present: + if isinstance(value, str) and len(value) > _MAX_VOCABULARY_VALUE_CHARS: + return None + try: + if value in distinct: + continue + distinct.add(value) + except TypeError: + return None # unhashable values (dicts / lists) have no cardinality signal + retained_bytes += len(value) if isinstance(value, str) else _NON_STRING_VALUE_BYTES + if len(distinct) > _MAX_VOCABULARY_VALUES or retained_bytes > _MAX_VOCABULARY_BYTES: + return None return CategoricalStats(distinct_count=len(distinct)) diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index a5841cddae..95dfb999f8 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -3,7 +3,14 @@ """Tests for per-column statistics.""" -from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats, quote_enumerations +from nemo_datasets_plugin.profiler.stats import ( + _MAX_VOCABULARY_BYTES, + _MAX_VOCABULARY_VALUE_CHARS, + _MAX_VOCABULARY_VALUES, + derive_probes, + derive_stats, + quote_enumerations, +) from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema @@ -33,15 +40,42 @@ def test_text_quality_flags_repetition_and_non_ascii(): assert stats.quality.non_ascii_ratio > 0.0 # accented characters -def test_cardinality_counts_are_always_stored(): - # distinct_count is a count, not row content, so it is always safe -- and it is the id-like - # signal the contract documents. - free_text = derive_stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)])) - assert free_text["t"].categorical.distinct_count == 50 # ~= row count -> id-like - +def test_cardinality_is_counted_while_the_column_is_a_vocabulary(): labels = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"])) assert labels["c"].categorical.distinct_count == 2 + # Still counted well past the point where every value is distinct: it is size, not repetition, + # that decides whether a column is a vocabulary. + many = derive_stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)])) + assert many["t"].categorical.distinct_count == 50 + + +def test_cardinality_stops_at_too_many_values(): + over = derive_stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES + 1)])) + assert over["t"].categorical is None # absence is the claim: not a vocabulary + assert over["t"].text is not None # ...but the column is still measured + + at_cap = derive_stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES)])) + assert at_cap["t"].categorical.distinct_count == _MAX_VOCABULARY_VALUES + + +def test_one_long_value_settles_it_without_counting(): + # The rule that does the real work: a vocabulary member is short by nature, so a single long + # value proves the column is not one -- on sight, rather than after a thousand of them. + values = ["yes", "no", "x" * (_MAX_VOCABULARY_VALUE_CHARS + 1)] + assert derive_stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None + + still_short = ["yes", "no", "x" * _MAX_VOCABULARY_VALUE_CHARS] + assert derive_stats([_feature("t", "string")], _rows("t", still_short))["t"].categorical.distinct_count == 3 + + +def test_cardinality_stops_on_total_bytes_before_the_count(): + # Values individually short enough and few enough, but heavy in aggregate. Without this bound + # the other two would admit 1024 x 256 chars -- four times the byte budget. + values = [f"{i:04d}" + "x" * 200 for i in range(_MAX_VOCABULARY_BYTES // 200)] + assert len(values) < _MAX_VOCABULARY_VALUES # the count bound is not what stops this + assert derive_stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None + def test_derive_stats_never_quotes_values(): # Quoting needs a role, and roles are not assigned when stats are measured. Filling them in From 4036b5a7ca21137ba935cbc1b58890c7ffb9731d Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 17:00:26 -0400 Subject: [PATCH 35/44] feat(datasets): bound and cheapen the per-character quality scan Phase 2 of the streaming spec. `TextQuality`'s three ratios become estimates over a bounded sample, and the two that can be computed in C without changing their meaning now are. These three are all the per-character work there is. Measured against the content probes, which scan the same strings: TextQuality 1.459s, all six probe regexes 0.039s -- 37x, because the probes have literal prefixes and these are character classes. Every other statistic is O(1) per row. So this is the whole reason reading a large dataset exhaustively is expensive, and bounding it is what makes Phase 5's "read every row" affordable. Sampling. `_QUALITY_SAMPLE_ROWS = 50_000`, evenly strided. Strided rather than random because two runs over the same bytes must agree -- randomness is precisely what `SamplingInfo.seed` existed to make reproducible, and that field was deleted on the grounds that the profiler makes no random choices. Strided rather than the head because shards arrive sorted often enough that the first rows of a large column are not a sample of it. Every denominator is the sample's own, so each ratio is unbiased rather than diluted. On accuracy, measured rather than asserted. Relative error is worst exactly where the value is nearest zero -- which is where there is nothing to flag. Injecting known corruption into a 20k-row column and sampling it three ways: at 30% corruption the worst estimate is off by 6.7%, at 80% by 0%, and the wild relative errors are all at 0.1% corruption where the absolute values are ~0.001. These signals exist to catch a column that is badly corrupt, and they are precise exactly when there is corruption to catch. On stride aliasing, also measured. A stride can lock onto periodic data, and HelpSteer2 is periodic -- two rated responses per prompt. Sampling its two phases separately: `prompt` is identical (both phases carry the same prompt) and `response` agrees to 0.35% on whitespace, differing at most 8% on the two rare-event ratios whose values there are 0.0003 and 0.0025. Inside the band these estimates already carry near zero, so no block-sampling scheme is bought. The hybrids. `str.count` over six ascii literals is not `\s` -- it misses U+00A0 and the rest of Unicode's spaces. `len(encode) - len` is not a character count -- it counts bytes of overhead, so one 4-byte emoji reports 3. Both are faster while measuring something else, and I had both in the spec at "10x" and "23x" before checking. What is exact *and* fast is branching on `str.isascii`, a C-level scan, and taking the shortcut only where it is provably identical. 1.6x on TextQuality exhaustively, 1.36x end-to-end on the demo, bit-identical ratios. Eight parametrized equivalence cases cover the divergences: U+00A0, ideographic and line/paragraph separators, a 4-byte codepoint, and mixed strings. Replacing either hybrid with its naive form fails four of them. Contract: `TextQuality` documents the three as estimates and says why. `stats_complete` gains a carve-out naming them, since it claims "proven facts, not estimates" -- the two only diverge on an unbounded read of a column past the bound. The carve-out goes away in Phase 5 when that field becomes `rows_complete`. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 19 ++++- .../nemo_datasets_plugin/profiler/stats.py | 70 +++++++++++++++++-- plugins/nemo-datasets/tests/test_stats.py | 57 +++++++++++++++ 3 files changed, 138 insertions(+), 8 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 9170ee5cff..cfc42110c4 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -181,8 +181,18 @@ class NumericStats(BaseModel): class TextQuality(BaseModel): - """Cheap, single-pass corruption signals for a text column. Flags training-wrecking data, not - toxicity / PII. + """Corruption signals for a text column. Flags training-wrecking data, not toxicity / PII. + + **Estimates, not counts**, and the only measurements in the profile that are. These three are + all the per-character work there is — every other statistic is O(1) per row, and the content + probes are literal searches costing a fraction of these — so scanning every row of a large + column costs more than the entire rest of the profile. They are also ratios, which a sample of + tens of thousands of rows pins down far past the precision anyone reads them to. Bounding them + is what makes reading every row of a dataset affordable. + + The sample is evenly strided: deterministic, so two runs over the same bytes agree, and spread + across the column rather than taken from its head, so a sorted shard does not decide the answer. + A column smaller than the bound is measured in full. """ whitespace_ratio: float = Field(ge=0.0, le=1.0, description="Padding / bad scraping.") @@ -448,7 +458,10 @@ class PartitionProfile(BaseModel): "True => `features`, `stats` and `classification` were computed over every row of every " "file in THIS partition: proven facts, not estimates. Only then can a consumer assert " "enum / required in a bridged JSON Schema, or read a verifiability coverage of 1.0 as " - "literal. Scoped to the partition because that is where it is decided — a corrupt shard " + "literal. `TextQuality` is the one exception and is always an estimate — it is bounded " + "independently of how much was read, for cost reasons its own docstring gives — though " + "the two only diverge on an unbounded read of a column larger than that bound. " + "Scoped to the partition because that is where it is decided — a corrupt shard " "in one partition says nothing about the measurements in another, and a fileset-wide " "flag quietly downgraded every partition to the worst one." ), diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index e44f763479..acd39000ca 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -212,22 +212,82 @@ def _cardinality(present: list[Any]) -> CategoricalStats | None: # passes over every character of every string and dominated total profiling time. _REPEAT_RUN = re.compile(r"(.)\1{3,}", re.DOTALL) +# What `\s` matches within ASCII, in the order `str.count` will be asked for them. +_ASCII_WHITESPACE = " \t\n\r\f\v" + +# Rows a column's quality ratios are measured over. These three are the only per-character work left +# in the profiler -- measured at 37x the cost of every content probe combined, and roughly fifteen +# times everything else in a column's measurement put together -- while every other statistic is +# O(1) per row. They are also ratios, which a sample of tens of thousands of rows pins down far past +# the precision anyone reads them to. Bounding them is what makes reading every row affordable. +_QUALITY_SAMPLE_ROWS = 50_000 + + +def _quality_sample(strings: list[str]) -> list[str]: + """The rows to measure quality over: all of them, or an evenly strided subset. + + Strided rather than random, because two runs over the same bytes must agree. Randomness is what + ``SamplingInfo.seed`` existed to make reproducible, and that field was deleted on the grounds + that the profiler makes no random choices and a seed would be theatre -- which should stay true. + + Strided rather than the head, because shards arrive sorted often enough that the first rows of a + large column are not a sample of it. A stride costs the same, needs no state, and spreads the + sample across the whole column. + + A stride can in principle alias against periodic data: a set with two rows per prompt, sampled + at stride two, sees one phase of every pair. Measured on exactly that shape -- HelpSteer2 rates + two responses per prompt -- the two phases agree to 0.35% on `whitespace_ratio` and differ by at + most 8% on the other two, whose values there are 0.0003 and 0.0025. That is well inside the band + these estimates already carry near zero, and does not buy a block-sampling scheme to avoid. + """ + if len(strings) <= _QUALITY_SAMPLE_ROWS: + return strings + return strings[:: len(strings) // _QUALITY_SAMPLE_ROWS] + + +def _whitespace_count(text: str) -> int: + """Whitespace characters, matching ``\\s`` exactly. + + The ASCII branch is not merely faster, it is the only one that may take the shortcut: within + ASCII ``\\s`` is precisely :data:`_ASCII_WHITESPACE`, so counting those six literals in C is the + same measurement. Outside it, ``\\s`` also matches U+00A0 and the rest of Unicode's spaces, which + the literal count would silently miss -- so the regex is a correctness fallback, not a slow path + kept for tidiness. + """ + if text.isascii(): + return sum(text.count(char) for char in _ASCII_WHITESPACE) + return _count_matches(_WHITESPACE_RUN, text) + + +def _non_ascii_count(text: str) -> int: + """Characters outside ASCII. ``str.isascii`` settles the common case in C without a scan. + + Deliberately not ``len(text.encode()) - len(text)``, which is faster still and answers a + different question: that counts *bytes* of encoding overhead, so a three-byte codepoint would + contribute two where this contributes one. + """ + if text.isascii(): + return 0 + return _count_matches(_NON_ASCII_RUN, text) + def _text_quality(strings: list[str]) -> TextQuality: + sample = _quality_sample(strings) total_chars = 0 whitespace = 0 non_ascii = 0 repetition_sum = 0.0 - for value in strings: + for value in sample: total_chars += len(value) - # str.count-style scanning in C rather than a per-character generator in Python. - whitespace += _count_matches(_WHITESPACE_RUN, value) - non_ascii += _count_matches(_NON_ASCII_RUN, value) + whitespace += _whitespace_count(value) + non_ascii += _non_ascii_count(value) repetition_sum += _repetition_score(value) return TextQuality( whitespace_ratio=whitespace / total_chars if total_chars else 0.0, non_ascii_ratio=non_ascii / total_chars if total_chars else 0.0, - repetition_score=repetition_sum / len(strings) if strings else 0.0, + # Every denominator is the sample's own, never the column's: each ratio is an estimate over + # the rows that were actually scanned, which is what keeps it unbiased rather than diluted. + repetition_score=repetition_sum / len(sample) if sample else 0.0, ) diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 95dfb999f8..5de424cb17 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -3,10 +3,18 @@ """Tests for per-column statistics.""" +import pytest +from nemo_datasets_plugin.profiler import stats from nemo_datasets_plugin.profiler.stats import ( _MAX_VOCABULARY_BYTES, _MAX_VOCABULARY_VALUE_CHARS, _MAX_VOCABULARY_VALUES, + _NON_ASCII_RUN, + _QUALITY_SAMPLE_ROWS, + _WHITESPACE_RUN, + _non_ascii_count, + _quality_sample, + _whitespace_count, derive_probes, derive_stats, quote_enumerations, @@ -40,6 +48,55 @@ def test_text_quality_flags_repetition_and_non_ascii(): assert stats.quality.non_ascii_ratio > 0.0 # accented characters +@pytest.mark.parametrize( + "text", + [ + "", + "plain ascii", + "tabs\tand\nnewlines\r\f\v", + "café naïve", # non-ascii letters, ascii spaces + "a b", # NO-BREAK SPACE: \s matches it, counting six ascii literals would not + "　

", # ideographic space, line separator, paragraph separator + "\U0001f600 emoji beyond the BMP", # 4-byte codepoint: byte overhead != character count + "mixed   café\tend", + ], +) +def test_quality_fast_paths_are_the_same_measurement_as_the_regexes(text): + # The whole risk in this change. `str.count` over six ascii literals is not `\s`, and + # `len(encode) - len` is not a character count -- either would be faster while quietly + # measuring something else. The fast path is only allowed where it is provably identical. + assert _whitespace_count(text) == sum(1 for _ in _WHITESPACE_RUN.finditer(text)) + assert _non_ascii_count(text) == sum(1 for _ in _NON_ASCII_RUN.finditer(text)) + + +def test_quality_sample_is_bounded_strided_and_deterministic(): + under = [f"r{i}" for i in range(_QUALITY_SAMPLE_ROWS)] + assert _quality_sample(under) is under # nothing to sample: measured in full + + over = [f"r{i}" for i in range(_QUALITY_SAMPLE_ROWS * 3)] + sample = _quality_sample(over) + assert len(sample) <= _QUALITY_SAMPLE_ROWS + 1 + assert _quality_sample(over) == sample # no RNG, so no seed to record and no run-to-run drift + # Spans the column rather than its head: the last sampled row is near the end. + assert over.index(sample[-1]) >= len(over) - 3 + + +def test_quality_is_measured_across_the_column_not_its_head(monkeypatch): + # Corruption confined to the second half. A head sample would report a clean column; a stride + # sees it. This is why the sample is strided and not simply the first N rows. + monkeypatch.setattr(stats, "_QUALITY_SAMPLE_ROWS", 4) + values = ["ordinary sentence"] * 8 + ["aaaaaaaaaaaa"] * 8 + quality = derive_stats([_feature("t", "string")], _rows("t", values))["t"].quality + assert quality.repetition_score > 0.4 + + +def test_a_column_under_the_bound_is_measured_exactly(monkeypatch): + monkeypatch.setattr(stats, "_QUALITY_SAMPLE_ROWS", 100) + values = ["ordinary sentence"] * 9 + ["aaaaaaaaaaaa"] + quality = derive_stats([_feature("t", "string")], _rows("t", values))["t"].quality + assert quality.repetition_score == pytest.approx(0.1) # exactly one corrupt row in ten + + def test_cardinality_is_counted_while_the_column_is_a_vocabulary(): labels = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"])) assert labels["c"].categorical.distinct_count == 2 From fff3fad99bb36b76af4b3e14c116b6577304fd50 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 17:43:51 -0400 Subject: [PATCH 36/44] feat(datasets): isolate each column's measurement from its neighbours MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3, first step. §9 of the spec asks for the failure-isolation tests before the streaming refactor, and for a guard the profiler did not have: per column, not just per partition. Two guards now, at different widths. The wide one still wraps the whole measure stage and catches anything structural -- schema derivation, classification -- that no single column owns. The new narrow one wraps each column, so a value no detector anticipated costs that column its measurements and nothing else. It used to cost the partition every measurement it had: one odd `role` and a seven-column profile came back with `dataset_type: unknown` and empty stats. A failed column is reported as `error` evidence rather than left as a gap, because absence from `stats` is also the ordinary sparse case -- a struct column with nothing worth measuring looks identical otherwise. The errors are appended after classify runs, so its own reasoning still reads first and the failure lands as a caveat on the result rather than as part of the case for it. Statistics and probes are now measured together in `measure_columns`. They read the same values and extracting a column out of the rows costs more than either measurement, so the pair was doing that work twice. That leaves `derive_stats` with no production caller and a duplicate of the extract/dedup loop, so it is deleted; the tests that used it now go through a three-line local helper that asserts no column failed, which is what they actually mean. `derive_probes` stays -- classification calls it when handed no probes. Behaviour is unchanged, and that is checked rather than asserted: a test compares `measure_columns` against the two functions it replaces across every dtype the dispatch knows, and the same comparison run over the real demo shards (13,392 rows, 10 columns across two partitions) is identical on both. Tests written before the change, per §9: the two failure domains have to stay distinguishable, and folding the read and measure loops together is what would blur them. A read failure is a FileError and leaves classification intact; a measurement failure is classification evidence and produces no FileError; and one partition's failure does not touch another's. Pinned now, before Phase 3 moves the loops. One of those pins is a wart worth naming: `stats_complete` stays True through a measurement failure, because it speaks to rows read and every row *was* read. It reads oddly next to `dataset_type: unknown`. Phase 5's rename to `rows_complete` is what makes it honest. Also corrected a comment in `_column_stats` that still described distinct_count as "always safe to store" -- Phase 1 made that false. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/pipeline.py | 13 +- .../nemo_datasets_plugin/profiler/stats.py | 79 +++++++---- plugins/nemo-datasets/tests/test_pipeline.py | 79 ++++++++++- plugins/nemo-datasets/tests/test_stats.py | 133 +++++++++++++----- 4 files changed, 242 insertions(+), 62 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 5676663374..b736ff52c0 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -35,7 +35,7 @@ from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits -from nemo_datasets_plugin.profiler.stats import derive_probes, derive_stats, quote_enumerations +from nemo_datasets_plugin.profiler.stats import measure_columns, quote_enumerations from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, DatasetProfile, @@ -241,14 +241,17 @@ def _measure( try: declared = _unify_schemas(arrow_schemas) if all_declared else None features = derive_features(partition_rows, declared) - stats = derive_stats(features, partition_rows) - # Probes are measured over every column, independent of the roles classify is about to - # assign, so a content signal survives a column name the alias table does not know. - probes = derive_probes(features, partition_rows) + # Statistics and probes together, one column at a time and each isolated. Probes cover every + # column independent of the roles classify is about to assign, so a content signal survives + # a column name the alias table does not know. + stats, probes, column_errors = measure_columns(features, partition_rows) classification = classify(features, stats, partition_rows, probes=probes, column_roles=column_roles) # Last, because the roles classification assigns are what decide whether a column's values # may be quoted at all — cardinality only bounds how many. quote_enumerations(features, stats, partition_rows) + # After classify, so its own reasoning reads first and a column that could not be measured + # is reported as a caveat on the result rather than as part of the case for it. + classification.evidence.extend(column_errors) return features, stats, classification except Exception as exc: detail = f"could not measure this partition: {type(exc).__name__}: {exc}" diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index acd39000ca..19dba4b5df 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -3,15 +3,17 @@ """Per-column statistics and content probes. -Given a partition's features and its sampled rows, measure each top-level column according to its -dtype: length quantiles and corruption signals for text, min/max/mean for numbers, chat-shape -signals for messages, and cardinality for both. The result is sparse — a column with nothing worth -measuring is omitted. Row values themselves are never stored here at all; a small controlled -vocabulary is added afterwards by :func:`quote_enumerations`, which gates on the column's role. - -:func:`derive_probes` additionally reads each column's *content* — answer markers, embedded -transcripts — as plain per-column counts. Those are measurements, not interpretations: what they -mean is classification's job, and keeping the looking here is what stops a content signal from +Given a partition's features and its rows, :func:`measure_columns` measures each top-level column +according to its dtype: length quantiles and corruption signals for text, min/max/mean for numbers, +chat-shape signals for messages, and a bounded vocabulary where the column has one. The result is +sparse — a column with nothing worth measuring is omitted — and each column is isolated, so one the +detectors cannot handle costs only itself. Row values themselves are never stored here at all; a +small controlled vocabulary is added afterwards by :func:`quote_enumerations`, which gates on role. + +The same pass reads each column's *content* — answer markers, embedded transcripts — as plain +per-column counts (:class:`ColumnProbes`, also reachable alone via :func:`derive_probes`, which +classification uses when it was handed no probes). Those are measurements, not interpretations: what +they mean is classification's job, and keeping the looking here is what stops a content signal from being reachable only through a correctly named column. """ @@ -25,6 +27,7 @@ from nemo_platform_plugin.files.dataset_profile import ( CategoricalStats, ColumnStats, + Evidence, FeatureSchema, MessageStats, NumericStats, @@ -57,23 +60,53 @@ _QUOTABLE_ROLES = frozenset({"label", "provenance", "meta", "rank"}) -def derive_stats(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> dict[str, ColumnStats]: - """Measure each top-level column. Keys are a subset of the feature names (sparse). +def measure_columns( + features: list[FeatureSchema], rows: list[dict[str, Any]] +) -> tuple[dict[str, ColumnStats], dict[str, ColumnProbes], list[Evidence]]: + """Measure every top-level column: its statistics and its content probes, in one pass each. - Never fills in ``categorical.values``: that needs the roles, which classification has not - assigned yet. :func:`quote_enumerations` adds them afterwards. + Each column is isolated. A value no detector anticipated -- a chat message whose ``role`` is a + number, a float where a string was declared -- costs that column its measurements and nothing + else, where previously it cost the partition every measurement it had. The failure is reported + as an ``error`` evidence rather than left as a silent gap, because a column absent from ``stats`` + is otherwise indistinguishable from one that simply had nothing worth measuring. + + This is the narrow half of the two guards the profiler runs. The wide one still wraps the whole + measure stage, and still catches anything structural -- schema derivation, classification -- that + is not attributable to a single column. + + Statistics and probes are measured together because they read the same values, and extracting a + column out of the rows costs more than either measurement. Neither fills in + ``categorical.values``: that needs the roles, which classification has not assigned yet, so + :func:`quote_enumerations` adds them afterwards. """ - total = len(rows) stats: dict[str, ColumnStats] = {} + probes: dict[str, ColumnProbes] = {} + errors: list[Evidence] = [] + total = len(rows) for feature in features: - # Parquet permits duplicate field names, and stats is keyed by name. Measuring the first and - # skipping the rest makes which one wins deterministic instead of "whichever came last". - if feature.name in stats: + # Parquet permits duplicate field names, and both maps are keyed by name. Measuring the + # first and skipping the rest makes which one wins deterministic instead of "whichever came + # last", and keeps stats and probes agreeing on the same one. + if feature.name in probes: + continue + values = [row.get(feature.name) for row in rows] + try: + column = _column_stats(feature, values, total) + probes[feature.name] = _column_probes(values) + except Exception as exc: + errors.append( + Evidence( + kind="error", + detail=( + f"column {feature.name!r} ({feature.dtype}) could not be measured: {type(exc).__name__}: {exc}" + ), + ) + ) continue - column = _column_stats(feature, [row.get(feature.name) for row in rows], total) if column is not None: stats[feature.name] = column - return stats + return stats, probes, errors def quote_enumerations( @@ -112,10 +145,8 @@ def _column_stats(feature: FeatureSchema, values: list[Any], total: int) -> Colu if strings: text = TextStats(chars=_quantiles([len(value) for value in strings])) quality = _text_quality(strings) - # distinct_count is always safe to store and is the id-like signal (~= rows_scanned) the - # contract documents. Only the values themselves are row data, and those are added later, - # by role. Withholding the count for high-cardinality strings dropped the signal precisely - # where it carries the most information. + # Only while the column looks like a controlled vocabulary; see `_cardinality`. The values + # themselves are row data and are added later, by role, never by size. categorical = _cardinality(present) elif feature.dtype == "bool": # The column that decides unpaired_preference deserves a measured class balance rather than @@ -434,7 +465,7 @@ def derive_probes(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> """Run the content probes over every top-level column, keyed by column name.""" probes: dict[str, ColumnProbes] = {} for feature in features: - # Duplicate parquet field names: first wins, matching derive_stats so the two agree on which. + # Duplicate parquet field names: first wins, matching `measure_columns` so both agree on which. if feature.name in probes: continue probes[feature.name] = _column_probes([row.get(feature.name) for row in rows]) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 692479c84d..a2f3684f64 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -529,7 +529,7 @@ def test_profile_degrades_one_partition_when_measurement_fails(tmp_path, monkeyp from nemo_datasets_plugin.profiler import pipeline as pipeline_module _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) - monkeypatch.setattr(pipeline_module, "derive_stats", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + monkeypatch.setattr(pipeline_module, "measure_columns", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise @@ -541,6 +541,83 @@ def test_profile_degrades_one_partition_when_measurement_fails(tmp_path, monkeyp assert "RuntimeError" in partition.classification.evidence[0].detail # says what failed +def test_a_read_failure_does_not_look_like_a_measurement_failure(tmp_path): + # The two failure domains have to stay distinguishable: a bad *file* is a FileError, and the + # rows that were readable still measure and classify normally. Folding the read and measure + # loops together is what would blur this, so it is pinned before that happens. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "q", "completion": "a"}]) + (tmp_path / "train-00001-of-00002.parquet").write_bytes(b"not parquet") + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + part = result.partitions[0] + assert [e.path for e in result.file_errors] == ["train-00001-of-00002.parquet"] + assert part.classification.dataset_type == "prompt_completion" # the readable rows still classify + assert "error" not in {e.kind for e in part.classification.evidence} + assert part.stats # ...and are still measured + + +def test_a_measurement_failure_does_not_look_like_a_read_failure(tmp_path, monkeypatch): + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) + monkeypatch.setattr(pipeline_module, "measure_columns", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.file_errors == [] # the file was fine; the data was odd + assert [e.kind for e in result.partitions[0].classification.evidence] == ["error"] + # `stats_complete` speaks to rows read, and every row *was* read -- so it stays True even though + # there are no stats. Pinned as it stands; the field means what it says once Phase 5 renames it. + assert result.partitions[0].stats_complete is True + + +def test_one_unmeasurable_column_does_not_cost_the_partition_its_classification(tmp_path, monkeypatch): + # The narrow guard, end to end. The column's failure reaches the profile as evidence, and + # everything the partition could still establish -- the other column's stats, the roles, the + # dataset type -- survives it. + from nemo_datasets_plugin.profiler import stats as stats_module + + real_column_stats = stats_module._column_stats + + def explode_on_completion(feature, values, total): + if feature.name == "completion": + raise RuntimeError("boom") + return real_column_stats(feature, values, total) + + monkeypatch.setattr(stats_module, "_column_stats", explode_on_completion) + _write_parquet(tmp_path / "train.parquet", [{"prompt": "q", "completion": "a"}]) + + part = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert "prompt" in part.stats and "completion" not in part.stats + assert part.classification.dataset_type == "prompt_completion" # roles come from names, not stats + assert any(e.kind == "error" and "'completion'" in e.detail for e in part.classification.evidence) + # The reasoning for the classification still reads first; the failure is a caveat on it. + assert part.classification.evidence[0].kind != "error" + + +def test_a_measurement_failure_is_scoped_to_its_own_partition(tmp_path, monkeypatch): + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + real_measure_columns = pipeline_module.measure_columns + + def poison_one_partition(features, rows): + if any(feature.name == "poison" for feature in features): + raise RuntimeError("boom") + return real_measure_columns(features, rows) + + _write_parquet(tmp_path / "good" / "train.parquet", [{"prompt": "q", "completion": "a"}]) + _write_parquet(tmp_path / "bad" / "train.parquet", [{"poison": 1}]) + monkeypatch.setattr(pipeline_module, "measure_columns", poison_one_partition) + + partitions = {p.name: p for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions} + + assert partitions["bad"].classification.dataset_type == "unknown" + assert partitions["good"].classification.dataset_type == "prompt_completion" + assert partitions["good"].stats # a neighbour's bad data costs this partition nothing + + def test_profile_is_deterministic(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) source = LocalFileSource(tmp_path) diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 5de424cb17..e0f8dd9a3a 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -4,7 +4,7 @@ """Tests for per-column statistics.""" import pytest -from nemo_datasets_plugin.profiler import stats +from nemo_datasets_plugin.profiler import stats as stats_module from nemo_datasets_plugin.profiler.stats import ( _MAX_VOCABULARY_BYTES, _MAX_VOCABULARY_VALUE_CHARS, @@ -16,12 +16,20 @@ _quality_sample, _whitespace_count, derive_probes, - derive_stats, + measure_columns, quote_enumerations, ) from nemo_platform_plugin.files.dataset_profile import ColumnStats, FeatureSchema +def _stats(features, rows): + """Statistics only. Asserts nothing failed: these tests measure values, not the guard, and a + swallowed exception would surface here as a confusing KeyError instead of its own message.""" + measured, _, errors = measure_columns(features, rows) + assert not errors, errors + return measured + + def _feature(name, dtype): return FeatureSchema(name=name, dtype=dtype) @@ -35,7 +43,7 @@ def _rows(name, values): def test_text_stats_length_quantiles_and_quality(): values = ["a", "bb", "ccc", "dddd"] - stats = derive_stats([_feature("t", "string")], _rows("t", values))["t"] + stats = _stats([_feature("t", "string")], _rows("t", values))["t"] assert stats.text.chars.max == 4 assert stats.text.chars.p50 in {2, 3} # nearest-rank over 4 values assert stats.quality is not None @@ -43,11 +51,72 @@ def test_text_stats_length_quantiles_and_quality(): def test_text_quality_flags_repetition_and_non_ascii(): - stats = derive_stats([_feature("t", "string")], _rows("t", ["aaaaaaaa", "héllo wörld"]))["t"] + stats = _stats([_feature("t", "string")], _rows("t", ["aaaaaaaa", "héllo wörld"]))["t"] assert stats.quality.repetition_score > 0.0 # the "aaaaaaaa" run assert stats.quality.non_ascii_ratio > 0.0 # accented characters +def test_one_bad_column_costs_only_itself(monkeypatch): + # The narrow guard. A value no detector anticipated used to cost the partition every measurement + # it had; it now costs its own column, and says so rather than leaving a silent gap. + real_column_stats = stats_module._column_stats + + def explode_on_one(feature, values, total): + if feature.name == "bad": + raise RuntimeError("boom") + return real_column_stats(feature, values, total) + + monkeypatch.setattr(stats_module, "_column_stats", explode_on_one) + + features = [_feature("good", "string"), _feature("bad", "string")] + measured, probes, errors = measure_columns(features, [{"good": "x", "bad": "y"}]) + + assert "good" in measured and "bad" not in measured + assert "good" in probes and "bad" not in probes # probes go with the column that failed + assert [e.kind for e in errors] == ["error"] + assert "'bad'" in errors[0].detail and "RuntimeError" in errors[0].detail + + +def test_measure_columns_agrees_with_the_unguarded_pair(): + # `measure_columns` is a refactor, not a change of measurement: it must produce exactly what the + # two functions it replaced produced, for every dtype the dispatch knows. Both paths are still + # here, so this is checkable directly rather than by reading. + features = [ + _feature("text", "string"), + _feature("count", "int64"), + _feature("score", "float64"), + _feature("flag", "bool"), + _feature("chat", "messages"), + _feature("meta", "struct"), + _feature("missing", "string"), + ] + rows = [ + { + "text": "a prompt ending in #### 42", + "count": i, + "score": i / 3, + "flag": i % 2 == 0, + "chat": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "there"}], + "meta": {"src": "x"}, + "missing": None, + } + for i in range(5) + ] + measured, probes, errors = measure_columns(features, rows) + + assert measured == _stats(features, rows) + assert probes == derive_probes(features, rows) + assert errors == [] + + +def test_a_column_with_nothing_to_measure_is_not_reported_as_an_error(): + # Absence from `stats` is the normal sparse case. Only a *failure* earns an error, or the two + # would be indistinguishable and the guard would cry wolf on every well-formed struct column. + measured, probes, errors = measure_columns([_feature("s", "struct")], [{"s": {"a": 1}}]) + assert measured == {} and errors == [] + assert "s" in probes + + @pytest.mark.parametrize( "text", [ @@ -84,35 +153,35 @@ def test_quality_sample_is_bounded_strided_and_deterministic(): def test_quality_is_measured_across_the_column_not_its_head(monkeypatch): # Corruption confined to the second half. A head sample would report a clean column; a stride # sees it. This is why the sample is strided and not simply the first N rows. - monkeypatch.setattr(stats, "_QUALITY_SAMPLE_ROWS", 4) + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 4) values = ["ordinary sentence"] * 8 + ["aaaaaaaaaaaa"] * 8 - quality = derive_stats([_feature("t", "string")], _rows("t", values))["t"].quality + quality = _stats([_feature("t", "string")], _rows("t", values))["t"].quality assert quality.repetition_score > 0.4 def test_a_column_under_the_bound_is_measured_exactly(monkeypatch): - monkeypatch.setattr(stats, "_QUALITY_SAMPLE_ROWS", 100) + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 100) values = ["ordinary sentence"] * 9 + ["aaaaaaaaaaaa"] - quality = derive_stats([_feature("t", "string")], _rows("t", values))["t"].quality + quality = _stats([_feature("t", "string")], _rows("t", values))["t"].quality assert quality.repetition_score == pytest.approx(0.1) # exactly one corrupt row in ten def test_cardinality_is_counted_while_the_column_is_a_vocabulary(): - labels = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"])) + labels = _stats([_feature("c", "string")], _rows("c", ["yes", "no", "yes", "no"])) assert labels["c"].categorical.distinct_count == 2 # Still counted well past the point where every value is distinct: it is size, not repetition, # that decides whether a column is a vocabulary. - many = derive_stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)])) + many = _stats([_feature("t", "string")], _rows("t", [f"unique-{i}" for i in range(50)])) assert many["t"].categorical.distinct_count == 50 def test_cardinality_stops_at_too_many_values(): - over = derive_stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES + 1)])) + over = _stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES + 1)])) assert over["t"].categorical is None # absence is the claim: not a vocabulary assert over["t"].text is not None # ...but the column is still measured - at_cap = derive_stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES)])) + at_cap = _stats([_feature("t", "string")], _rows("t", [f"v{i}" for i in range(_MAX_VOCABULARY_VALUES)])) assert at_cap["t"].categorical.distinct_count == _MAX_VOCABULARY_VALUES @@ -120,10 +189,10 @@ def test_one_long_value_settles_it_without_counting(): # The rule that does the real work: a vocabulary member is short by nature, so a single long # value proves the column is not one -- on sight, rather than after a thousand of them. values = ["yes", "no", "x" * (_MAX_VOCABULARY_VALUE_CHARS + 1)] - assert derive_stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None + assert _stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None still_short = ["yes", "no", "x" * _MAX_VOCABULARY_VALUE_CHARS] - assert derive_stats([_feature("t", "string")], _rows("t", still_short))["t"].categorical.distinct_count == 3 + assert _stats([_feature("t", "string")], _rows("t", still_short))["t"].categorical.distinct_count == 3 def test_cardinality_stops_on_total_bytes_before_the_count(): @@ -131,18 +200,18 @@ def test_cardinality_stops_on_total_bytes_before_the_count(): # the other two would admit 1024 x 256 chars -- four times the byte budget. values = [f"{i:04d}" + "x" * 200 for i in range(_MAX_VOCABULARY_BYTES // 200)] assert len(values) < _MAX_VOCABULARY_VALUES # the count bound is not what stops this - assert derive_stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None + assert _stats([_feature("t", "string")], _rows("t", values))["t"].categorical is None def test_derive_stats_never_quotes_values(): # Quoting needs a role, and roles are not assigned when stats are measured. Filling them in # afterwards rather than redacting means a skipped pass stores nothing instead of leaking. - stats = derive_stats([_feature("c", "string")], _rows("c", ["yes", "no"])) + stats = _stats([_feature("c", "string")], _rows("c", ["yes", "no"])) assert stats["c"].categorical.values is None def test_bool_column_gets_a_measured_class_balance(): - stats = derive_stats([_feature("label", "bool")], _rows("label", [True, False, True])) + stats = _stats([_feature("label", "bool")], _rows("label", [True, False, True])) assert stats["label"].categorical.distinct_count == 2 @@ -150,14 +219,14 @@ def test_bool_column_gets_a_measured_class_balance(): def test_numeric_stats_and_cardinality(): - stats = derive_stats([_feature("n", "int64")], _rows("n", [0, 4, 2, 2, 3]))["n"] + stats = _stats([_feature("n", "int64")], _rows("n", [0, 4, 2, 2, 3]))["n"] assert (stats.numeric.min, stats.numeric.max) == (0.0, 4.0) assert stats.numeric.mean == 2.2 assert stats.categorical.distinct_count == 4 # {0, 2, 3, 4} def test_numeric_cardinality_counts_without_quoting(): - stats = derive_stats([_feature("n", "int64")], _rows("n", [1, 2, 3]))["n"] + stats = _stats([_feature("n", "int64")], _rows("n", [1, 2, 3]))["n"] assert stats.categorical.distinct_count == 3 assert stats.categorical.values is None @@ -166,13 +235,13 @@ def test_numeric_stats_ignore_non_finite_values(): # NaN / +-inf poison min/max/mean and serialize to JSON null, which then fails to re-validate # against NumericStats' required floats -- making the whole profile unreadable. Drop them. values = [1.0, float("nan"), 3.0, float("inf"), float("-inf"), 5.0] - stats = derive_stats([_feature("n", "float64")], _rows("n", values))["n"] + stats = _stats([_feature("n", "float64")], _rows("n", values))["n"] assert (stats.numeric.min, stats.numeric.max, stats.numeric.mean) == (1.0, 5.0, 3.0) ColumnStats.model_validate_json(stats.model_dump_json()) # round-trips: no NaN/inf leaked into JSON def test_numeric_all_non_finite_yields_no_numeric_summary(): - stats = derive_stats([_feature("n", "float64")], _rows("n", [float("nan"), float("inf")])) + stats = _stats([_feature("n", "float64")], _rows("n", [float("nan"), float("inf")])) assert stats.get("n") is None or stats["n"].numeric is None @@ -184,7 +253,7 @@ def test_message_stats_shape_signals(): {"m": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello there"}]}, {"m": [{"role": "user", "content": "again"}, {"role": "assistant", "content": "yes"}]}, ] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.turns.max == 2 assert stats.messages.roles_seen == ["user", "assistant"] # first-seen order assert stats.messages.ends_with_assistant_rate == 1.0 @@ -194,20 +263,20 @@ def test_message_stats_shape_signals(): def test_message_stats_detects_tool_calls_and_user_ending(): rows = [{"m": [{"role": "user", "content": "run"}, {"role": "assistant", "tool_calls": [{"id": "1"}]}]}] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.has_tool_calls is True assert stats.messages.ends_with_assistant_rate == 1.0 # last turn is the assistant tool call def test_message_ends_with_user_turn_is_prompt_only_signal(): rows = [{"m": [{"role": "user", "content": "solve"}]}] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.ends_with_assistant_rate == 0.0 def test_message_stats_read_sharegpt_from_value(): rows = [{"m": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello there"}]}] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.roles_seen == ["human", "gpt"] # verbatim, not normalized assert stats.messages.content_chars.max == len("hi") + len("hello there") assert stats.messages.ends_with_assistant_rate == 1.0 # "gpt" is the responder turn @@ -217,7 +286,7 @@ def test_assistant_equivalent_roles_count_as_the_training_target(): # Matching only the literal "assistant" made every other convention look prompt-only. for responder in ("assistant", "gpt", "bot", "model", "AI"): rows = [{"m": [{"role": "user", "content": "q"}, {"role": responder, "content": "a"}]}] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.ends_with_assistant_rate == 1.0, responder @@ -225,7 +294,7 @@ def test_non_string_role_does_not_break_measurement(): # roles_seen is typed list[str]; a numeric role used to raise a ValidationError from inside the # one stage the pipeline did not guard, aborting the whole profile. rows = [{"m": [{"role": 1, "content": "hi"}]}] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.roles_seen == ["1"] @@ -233,14 +302,14 @@ def test_declared_but_unset_tool_calls_is_not_tool_use(): # parquet materializes every declared struct field, so `"tool_calls" in message` reported tool # use for any schema that merely declares the field. rows = [{"m": [{"role": "user", "content": "hi", "tool_calls": None}]}] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.has_tool_calls is False def test_message_content_parts_tolerate_non_string_text(): # A VLM-style content part whose "text" key is present but not a string must not crash measurement. rows = [{"m": [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": None}]}]}] - stats = derive_stats([_feature("m", "messages")], rows)["m"] + stats = _stats([_feature("m", "messages")], rows)["m"] assert stats.messages.content_chars.max == 0 # no measurable text, and no crash @@ -250,11 +319,11 @@ def test_message_content_parts_tolerate_non_string_text(): def test_unmeasured_dtypes_are_omitted(): features = [_feature("s", "struct"), _feature("j", "json")] rows = [{"s": {"a": 1}, "j": object()}] - assert derive_stats(features, rows) == {} + assert _stats(features, rows) == {} def test_null_rate_is_reported(): - stats = derive_stats([_feature("t", "string")], _rows("t", ["a", None, "c", None]))["t"] + stats = _stats([_feature("t", "string")], _rows("t", ["a", None, "c", None]))["t"] assert stats.null_rate == 0.5 @@ -314,7 +383,7 @@ def _quoted(name, dtype, values, role): feature = _feature(name, dtype) feature.semantic_role = role rows = _rows(name, values) - stats = derive_stats([feature], rows) + stats = _stats([feature], rows) quote_enumerations([feature], stats, rows) return stats[name].categorical.values From 7a5cded6b6b5a3ef79da79d81cda3174ca92dbdb Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 17:59:37 -0400 Subject: [PATCH 37/44] refactor(datasets): measure each column with an accumulator Phase 3, second step. `_column_stats` and `_column_probes` become a `ColumnAccumulator` per column, chosen once on dtype: `update` folds a batch in and keeps no reference to it, `finalize` turns what was folded into the stored blocks. The point is one property. A column measured in pieces has to measure the same as one measured whole, or batching would quietly change the numbers -- and the batch size is an implementation detail no reader of a profile could see. That is what lets the caller stop materialising a partition before measuring it, which is the whole of what Phase 3 is for. Tested directly: every dtype the dispatch knows, fed in 1, 2, 3 and 7 chunks, must finalize identically. Five accumulators. The base class is the entire measurement for a dtype with no statistics of its own -- a struct, a list -- because the probes run over every column whatever its type; that is also what `derive_probes` now uses, so probes alone cost the scan and nothing else. String, numeric, bool and messages add their state by overriding two methods. `_cardinality` becomes `_Vocabulary`, which is the piece that most wanted to be stateful: it was already a bounded scan with three early exits, and as an object it simply stops and drops what it held. Same three bounds, same results. Behaviour is unchanged, and checked rather than argued: HEAD's stats.py is loaded side by side with the new one and both are run over the real demo shards. Stats and probes are identical on both partitions, 13,392 rows across 10 columns. Two tests that verified the previous step had gone circular now that `_stats` in the test file *is* `measure_columns`, so they assert real content per dtype instead -- which columns get which blocks, and that a struct with no nulls is correctly absent while an all-null string column is kept for its null rate alone. Honest about what is not yet bounded: a string accumulator still retains its strings, because the quality stride needs the column's length to place its sample and that is not known until the last batch, and quantiles still sort every length. Both are the reservoir's and the parquet footers' job -- the next two steps -- and until then they cost exactly what materialising the column already cost. The module docstring says so rather than implying an O(1) that is not there yet. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/stats.py | 373 +++++++++++++----- plugins/nemo-datasets/tests/test_pipeline.py | 13 +- plugins/nemo-datasets/tests/test_stats.py | 79 +++- 3 files changed, 356 insertions(+), 109 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 19dba4b5df..1b38166a0c 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -15,6 +15,17 @@ classification uses when it was handed no probes). Those are measurements, not interpretations: what they mean is classification's job, and keeping the looking here is what stops a content signal from being reachable only through a correctly named column. + +The measuring itself is done by a :class:`ColumnAccumulator` per column, chosen once on dtype. An +accumulator folds batches in and keeps no reference to them, so a column measured in pieces gives +the same answer as one measured whole — the property that lets a caller stop materialising a +partition before it can measure it. The base class is the entire measurement for a dtype with no +statistics of its own, because the probes run over every column whatever its type. + +Two things here are still sized by the column rather than bounded: the retained strings a quality +stride needs to place its sample, and the lengths a quantile needs to sort. Both are what the +reservoir and the parquet footer row counts are for, and until then they cost what materialising the +column already cost. """ from __future__ import annotations @@ -83,17 +94,20 @@ def measure_columns( stats: dict[str, ColumnStats] = {} probes: dict[str, ColumnProbes] = {} errors: list[Evidence] = [] - total = len(rows) for feature in features: # Parquet permits duplicate field names, and both maps are keyed by name. Measuring the # first and skipping the rest makes which one wins deterministic instead of "whichever came # last", and keeps stats and probes agreeing on the same one. if feature.name in probes: continue - values = [row.get(feature.name) for row in rows] + accumulator = _accumulator_for(feature) try: - column = _column_stats(feature, values, total) - probes[feature.name] = _column_probes(values) + # One batch, being every row this partition holds. The accumulator does not care: fed in + # pieces it gives the same answer, which is what lets the caller stop materialising the + # partition first. + accumulator.update([row.get(feature.name) for row in rows]) + column, probe = accumulator.finalize() + probes[feature.name] = probe except Exception as exc: errors.append( Evidence( @@ -135,44 +149,261 @@ def quote_enumerations( column.categorical.values = sorted(str(value) for value in distinct) -def _column_stats(feature: FeatureSchema, values: list[Any], total: int) -> ColumnStats | None: - present = [value for value in values if value is not None] - null_rate = (total - len(present)) / total if total else 0.0 +class ColumnAccumulator: + """Measures one top-level column, over however many batches it is handed. + + ``update`` folds a batch in and keeps no reference to it; ``finalize`` turns what was folded into + the stored blocks. Splitting a column across calls gives the same answer as one call with all of + it, which is the property that lets the caller stop materialising a partition before measuring it. + + The base class is the whole measurement for a dtype with no statistics of its own — a struct, a + list, anything the dispatch does not recognise — because the content probes run over every column + regardless of type. Subclasses add their dtype's state by overriding ``_observe`` and ``_blocks``. + """ + + def __init__(self) -> None: + self.rows = 0 + self._nulls = 0 + self._non_empty = 0 + self._texts = 0 + self._extractable_answer = 0 + self._transcript_marker = 0 + + def update(self, values: list[Any]) -> None: + """Fold one batch of this column's values in, one entry per row.""" + present: list[Any] = [] + for value in values: + self.rows += 1 + if value is None: + self._nulls += 1 + else: + present.append(value) + if value not in ("", [], {}): + self._non_empty += 1 + text = _probe_text(value) + if text is not None: + self._texts += 1 + if _GSM8K_ANSWER.search(text) or _BOXED_ANSWER.search(text): + self._extractable_answer += 1 + if _TRANSCRIPT_MARKER.search(text): + self._transcript_marker += 1 + self._observe(present) + + def finalize(self) -> tuple[ColumnStats | None, ColumnProbes]: + """The column's stored measurements, and its probe counts. + + Stats are None when there was nothing worth measuring, which keeps the map sparse. Probes are + always returned: a column of nothing is a finding classification is entitled to read. + """ + blocks = self._blocks() + null_rate = self._nulls / self.rows if self.rows else 0.0 + column = ColumnStats(null_rate=null_rate, **blocks) + if not any(blocks.values()) and null_rate == 0.0: + column = None + return column, ColumnProbes( + rows=self.rows, + non_empty=self._non_empty, + texts=self._texts, + extractable_answer=self._extractable_answer, + transcript_marker=self._transcript_marker, + ) + + def _observe(self, present: list[Any]) -> None: + """Fold this batch's non-null values into the dtype's own state. The base column has none.""" + + def _blocks(self) -> dict[str, Any]: + """The dtype-specific ``ColumnStats`` blocks. The base column contributes none.""" + return {} + + +class _Vocabulary: + """Distinct values, for as long as the column still looks like a controlled vocabulary. + + Stops the moment it stops looking like one and drops what it had, which is the whole point: + counting distinct values exactly means *retaining* them, so on a free-text column this set grows + to hold the column. Today that costs little, because the rows are held anyway and the set stores + pointers into them -- 2.6 MB beside 61.4 MB of resident rows. It is the fold this is becoming + that makes it matter: once a batch is folded and discarded, this set is the *sole owner* of every + value it kept, and two text columns cost 46.8 MB against 0.163 MB for every other accumulator + combined. Unbounded, it is the one thing that would make the fold O(rows) again. + + Three bounds rather than one. A count alone bounds cardinality but not bytes, and 1024 reasoning + traces is 32 MB. The middle bound does the real work: it asks what the column *is* rather than + how many values it holds, in the same way the role gate on quoting does. A vocabulary member is + short by nature, so a single long value settles the question on sight, which is why free-text + columns stop here almost immediately instead of after 1024 values. + + The values themselves are never handed out here. They are row content, gated on role rather than + on size, and :func:`quote_enumerations` adds them once classification has assigned one. + """ + + def __init__(self) -> None: + self._values: set[Any] = set() + self._bytes = 0 + self._saturated = False + + def update(self, present: list[Any]) -> None: + if self._saturated: + return + for value in present: + if isinstance(value, str) and len(value) > _MAX_VOCABULARY_VALUE_CHARS: + return self._saturate() + try: + if value in self._values: + continue + self._values.add(value) + except TypeError: + return self._saturate() # unhashable values (dicts / lists) have no cardinality signal + self._bytes += len(value) if isinstance(value, str) else _NON_STRING_VALUE_BYTES + if len(self._values) > _MAX_VOCABULARY_VALUES or self._bytes > _MAX_VOCABULARY_BYTES: + return self._saturate() + + def _saturate(self) -> None: + self._values = set() # release what was held; holding it is the cost this bound exists to cap + self._saturated = True + + def finalize(self) -> CategoricalStats | None: + return None if self._saturated else CategoricalStats(distinct_count=len(self._values)) + + +class StringAccumulator(ColumnAccumulator): + """A ``string`` column: length quantiles, corruption ratios, and a vocabulary if it has one.""" + + def __init__(self) -> None: + super().__init__() + self._strings: list[str] = [] + self._vocabulary = _Vocabulary() + + def _observe(self, present: list[Any]) -> None: + # The strings are retained because the quality stride needs the column's length to place its + # sample, and that is not known until the last batch. Bounding this is what the reservoir and + # the footer row count are for; until then it costs what materialising the column already did. + self._strings.extend(value for value in present if isinstance(value, str)) + self._vocabulary.update(present) + + def _blocks(self) -> dict[str, Any]: + text = quality = None + if self._strings: + text = TextStats(chars=_quantiles([len(value) for value in self._strings])) + quality = _text_quality(self._strings) + return {"text": text, "quality": quality, "categorical": self._vocabulary.finalize()} + + +class NumericAccumulator(ColumnAccumulator): + """An ``int*`` / ``uint*`` / ``float*`` column: running extrema and mean, plus a vocabulary.""" + + def __init__(self) -> None: + super().__init__() + self._min = math.inf + self._max = -math.inf + self._sum = 0.0 + self._count = 0 + self._vocabulary = _Vocabulary() + + def _observe(self, present: list[Any]) -> None: + for value in present: + # Non-finite floats (NaN / +-inf) are dropped: they serialize to JSON null and then fail + # to re-validate against NumericStats' required floats, making the profile unreadable on + # its next load. bool is an int in Python and is not a number here. + if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value): + number = float(value) + self._min = min(self._min, number) + self._max = max(self._max, number) + self._sum += number + self._count += 1 + self._vocabulary.update(present) + + def _blocks(self) -> dict[str, Any]: + numeric = None + if self._count: + numeric = NumericStats(min=self._min, max=self._max, mean=self._sum / self._count) + return {"numeric": numeric, "categorical": self._vocabulary.finalize()} + + +class BoolAccumulator(ColumnAccumulator): + """A ``bool`` column. The column that decides unpaired_preference deserves a measured class + balance rather than no stats at all, and two values is a vocabulary by any reading.""" + + def __init__(self) -> None: + super().__init__() + self._vocabulary = _Vocabulary() + + def _observe(self, present: list[Any]) -> None: + self._vocabulary.update(present) + + def _blocks(self) -> dict[str, Any]: + return {"categorical": self._vocabulary.finalize()} + + +class MessageAccumulator(ColumnAccumulator): + """A ``messages`` column: turn and length distributions, the roles seen, and chat-shape rates.""" + + def __init__(self) -> None: + super().__init__() + self._conversations = 0 + self._turns: list[int] = [] + self._content_chars: list[int] = [] + self._roles_seen: list[str] = [] + self._ends_with_assistant = 0 + self._valid_alternation = 0 + self._has_tool_calls = False + + def _observe(self, present: list[Any]) -> None: + for messages in present: + if not isinstance(messages, list): + continue + self._conversations += 1 + self._turns.append(len(messages)) + total_content = 0 + for message in messages: + if not isinstance(message, dict): + continue + role = _role_of(message) + if role is not None: + # Coerced to str because roles_seen is typed list[str] and a non-string role + # would fail validation. Reported verbatim otherwise: the contract is explicit + # that an unexpected role is the finding worth surfacing, not something to + # normalize away. + role = role if isinstance(role, str) else str(role) + if role not in self._roles_seen: + self._roles_seen.append(role) + total_content += _content_len(_message_field(message, "content", "value")) + # `.get` truthiness, not `in`: parquet materializes every declared struct field, so a + # schema that merely declares tool_calls would otherwise report tool use on every row. + if message.get("tool_calls") or role == "tool": + self._has_tool_calls = True + self._content_chars.append(total_content) + if messages and isinstance(messages[-1], dict) and _is_assistant_role(_role_of(messages[-1])): + self._ends_with_assistant += 1 + if _valid_alternation(messages): + self._valid_alternation += 1 + + def _blocks(self) -> dict[str, Any]: + if not self._conversations: + return {"messages": None} + return { + "messages": MessageStats( + turns=_quantiles(self._turns), + content_chars=_quantiles(self._content_chars), + roles_seen=self._roles_seen, + ends_with_assistant_rate=self._ends_with_assistant / self._conversations, + valid_alternation_rate=self._valid_alternation / self._conversations, + has_tool_calls=self._has_tool_calls, + ) + } - text = numeric = messages = categorical = quality = None + +def _accumulator_for(feature: FeatureSchema) -> ColumnAccumulator: + """The accumulator that knows how to measure this column, dispatched once on its dtype.""" if feature.dtype == "string": - strings = [value for value in present if isinstance(value, str)] - if strings: - text = TextStats(chars=_quantiles([len(value) for value in strings])) - quality = _text_quality(strings) - # Only while the column looks like a controlled vocabulary; see `_cardinality`. The values - # themselves are row data and are added later, by role, never by size. - categorical = _cardinality(present) - elif feature.dtype == "bool": - # The column that decides unpaired_preference deserves a measured class balance rather than - # no stats at all. - categorical = _cardinality(present) - elif _is_numeric(feature.dtype): - # Drop non-finite floats (NaN / +-inf): they serialize to JSON null and then fail to - # re-validate against NumericStats' required floats, which would make the whole profile - # unreadable on the next load. - numbers = [ - float(value) - for value in present - if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) - ] - if numbers: - numeric = NumericStats(min=min(numbers), max=max(numbers), mean=sum(numbers) / len(numbers)) - categorical = _cardinality(present) - elif feature.dtype == "messages": - messages = _message_stats([value for value in present if isinstance(value, list)]) - - column = ColumnStats( - null_rate=null_rate, text=text, numeric=numeric, messages=messages, categorical=categorical, quality=quality - ) - if not any([text, numeric, messages, categorical, quality]) and null_rate == 0.0: - return None # nothing worth measuring - return column + return StringAccumulator() + if feature.dtype == "bool": + return BoolAccumulator() + if feature.dtype == "messages": + return MessageAccumulator() + if _is_numeric(feature.dtype): + return NumericAccumulator() + return ColumnAccumulator() def _is_numeric(dtype: str) -> bool: @@ -193,46 +424,6 @@ def at(percentile: int) -> int: return Quantiles(p50=at(50), p95=at(95), p99=at(99), max=ordered[-1] if ordered else 0) -def _cardinality(present: list[Any]) -> CategoricalStats | None: - """The distinct-value count, for as long as the column looks like a controlled vocabulary. - - Returns None the moment it stops looking like one, discarding what it had accumulated. Counting - distinct values exactly means *retaining* them, so on a free-text column this set grows to hold - the column, to report that `response` had 9,954 distinct values in 10,000 rows -- which is "this - is free text", which the role marker and the length quantiles already say for nothing. - - Today that costs little, because the rows are held anyway and the set stores pointers into them: - measured at 2.6 MB beside 61.4 MB of resident rows. The bound is here for the pipeline this is - becoming. Once a batch is folded and discarded, a distinct set is the *sole owner* of every value - it kept, and the same two columns cost 46.8 MB against 0.163 MB for every other accumulator - combined. Unbounded, this is the one thing that would make a streaming fold O(rows) again. - - Three bounds rather than one. A count alone bounds cardinality but not bytes, and 1024 reasoning - traces is 32 MB. The middle bound is the one that does the real work: it asks what the column - *is* rather than how many values it holds, in the same way the role gate on quoting does. A - vocabulary member is short by nature, so a single 32 KB value settles the question on sight, - which is why free-text columns exit here almost immediately instead of after 1024 values. - - The values themselves are still never returned here. They are row content, gated on role rather - than on size, and :func:`quote_enumerations` adds them once classification has assigned one. - """ - distinct: set[Any] = set() - retained_bytes = 0 - for value in present: - if isinstance(value, str) and len(value) > _MAX_VOCABULARY_VALUE_CHARS: - return None - try: - if value in distinct: - continue - distinct.add(value) - except TypeError: - return None # unhashable values (dicts / lists) have no cardinality signal - retained_bytes += len(value) if isinstance(value, str) else _NON_STRING_VALUE_BYTES - if len(distinct) > _MAX_VOCABULARY_VALUES or retained_bytes > _MAX_VOCABULARY_BYTES: - return None - return CategoricalStats(distinct_count=len(distinct)) - - # --- text quality -------------------------------------------------------------------------------- @@ -462,27 +653,23 @@ class ColumnProbes: def derive_probes(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> dict[str, ColumnProbes]: - """Run the content probes over every top-level column, keyed by column name.""" + """Run the content probes over every top-level column, keyed by column name. + + The probes alone, through the base accumulator, which carries no dtype state — so this costs the + scan and nothing else. Classification calls it when it was handed no probes of its own; + :func:`measure_columns` is the path that wants the statistics as well. + """ probes: dict[str, ColumnProbes] = {} for feature in features: # Duplicate parquet field names: first wins, matching `measure_columns` so both agree on which. if feature.name in probes: continue - probes[feature.name] = _column_probes([row.get(feature.name) for row in rows]) + accumulator = ColumnAccumulator() + accumulator.update([row.get(feature.name) for row in rows]) + probes[feature.name] = accumulator.finalize()[1] return probes -def _column_probes(values: list[Any]) -> ColumnProbes: - texts = [text for value in values if (text := _probe_text(value)) is not None] - return ColumnProbes( - rows=len(values), - non_empty=sum(1 for value in values if value not in (None, "", [], {})), - texts=len(texts), - extractable_answer=sum(1 for text in texts if _GSM8K_ANSWER.search(text) or _BOXED_ANSWER.search(text)), - transcript_marker=sum(1 for text in texts if _TRANSCRIPT_MARKER.search(text)), - ) - - def _probe_text(value: Any) -> str | None: """The text a probe reads from one cell: the string itself, or a chat column's final turn. diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index a2f3684f64..648ba4efb9 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -578,14 +578,17 @@ def test_one_unmeasurable_column_does_not_cost_the_partition_its_classification( # dataset type -- survives it. from nemo_datasets_plugin.profiler import stats as stats_module - real_column_stats = stats_module._column_stats + real_accumulator_for = stats_module._accumulator_for - def explode_on_completion(feature, values, total): - if feature.name == "completion": + class Boom(stats_module.ColumnAccumulator): + def _observe(self, present): raise RuntimeError("boom") - return real_column_stats(feature, values, total) - monkeypatch.setattr(stats_module, "_column_stats", explode_on_completion) + monkeypatch.setattr( + stats_module, + "_accumulator_for", + lambda feature: Boom() if feature.name == "completion" else real_accumulator_for(feature), + ) _write_parquet(tmp_path / "train.parquet", [{"prompt": "q", "completion": "a"}]) part = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index e0f8dd9a3a..ec8cc60c62 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -59,14 +59,17 @@ def test_text_quality_flags_repetition_and_non_ascii(): def test_one_bad_column_costs_only_itself(monkeypatch): # The narrow guard. A value no detector anticipated used to cost the partition every measurement # it had; it now costs its own column, and says so rather than leaving a silent gap. - real_column_stats = stats_module._column_stats + real_accumulator_for = stats_module._accumulator_for - def explode_on_one(feature, values, total): - if feature.name == "bad": + class Boom(stats_module.ColumnAccumulator): + def _observe(self, present): raise RuntimeError("boom") - return real_column_stats(feature, values, total) - monkeypatch.setattr(stats_module, "_column_stats", explode_on_one) + monkeypatch.setattr( + stats_module, + "_accumulator_for", + lambda feature: Boom() if feature.name == "bad" else real_accumulator_for(feature), + ) features = [_feature("good", "string"), _feature("bad", "string")] measured, probes, errors = measure_columns(features, [{"good": "x", "bad": "y"}]) @@ -77,10 +80,57 @@ def explode_on_one(feature, values, total): assert "'bad'" in errors[0].detail and "RuntimeError" in errors[0].detail -def test_measure_columns_agrees_with_the_unguarded_pair(): - # `measure_columns` is a refactor, not a change of measurement: it must produce exactly what the - # two functions it replaced produced, for every dtype the dispatch knows. Both paths are still - # here, so this is checkable directly rather than by reading. +# One column per dtype the dispatch knows, each carrying the awkward cases: nulls, empties, a +# non-finite float, a value long enough to saturate a vocabulary, both chat spellings. +_DTYPE_VALUES = { + "string": ["a prompt #### 42", "héllo wörld", "", "aaaaaaaa", None, "x" * 300, "yes", "yes"], + "int64": [1, 2, 2, None, 3, -5], + "float64": [1.5, float("nan"), 2.5, None, float("inf"), 0.0], + "bool": [True, False, True, None], + "messages": [ + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "there"}], + [{"from": "human", "value": "q"}, {"from": "gpt", "value": "\\boxed{4}"}], + [{"role": "user", "content": "x", "tool_calls": [{"n": 1}]}], + None, + [], + ], + "struct": [{"a": 1}, None, {"b": 2}], +} + + +@pytest.mark.parametrize("dtype", sorted(_DTYPE_VALUES)) +@pytest.mark.parametrize("chunks", [1, 2, 3, 7]) +def test_an_accumulator_folds_rather_than_buffers(dtype, chunks): + # The property the fold rests on: a column split across calls has to measure the same as one + # handed over whole. Without it, batching would quietly change the numbers -- and the batch size + # is an implementation detail no reader of the profile could see. + values = _DTYPE_VALUES[dtype] * 5 + feature = _feature("c", dtype) + + whole = stats_module._accumulator_for(feature) + whole.update(values) + + in_pieces = stats_module._accumulator_for(feature) + step = max(1, -(-len(values) // chunks)) + for start in range(0, len(values), step): + in_pieces.update(values[start : start + step]) + + assert in_pieces.finalize() == whole.finalize() + + +def test_the_typed_accumulators_probe_exactly_as_the_bare_one_does(): + # Probe counting lives on the base class, so every dtype gets it for free. If a subclass ever + # shadows that, a chat column would quietly stop contributing verifiability evidence. + features = [_feature(dtype, dtype) for dtype in sorted(_DTYPE_VALUES)] + rows = [dict(zip(sorted(_DTYPE_VALUES), values)) for values in zip(*_DTYPE_VALUES.values())] + + _, probes, errors = measure_columns(features, rows) + + assert errors == [] + assert probes == derive_probes(features, rows) # derive_probes goes through the base class + + +def test_measure_columns_measures_every_dtype_the_dispatch_knows(): features = [ _feature("text", "string"), _feature("count", "int64"), @@ -104,9 +154,16 @@ def test_measure_columns_agrees_with_the_unguarded_pair(): ] measured, probes, errors = measure_columns(features, rows) - assert measured == _stats(features, rows) - assert probes == derive_probes(features, rows) assert errors == [] + assert measured["text"].text is not None and measured["text"].quality is not None + assert measured["count"].numeric.min == 0.0 + assert measured["count"].categorical.distinct_count == 5 + assert measured["score"].numeric.mean == pytest.approx(sum(i / 3 for i in range(5)) / 5) + assert measured["flag"].categorical.distinct_count == 2 + assert measured["chat"].messages.roles_seen == ["user", "assistant"] + assert "meta" not in measured # a struct with no nulls has nothing worth measuring + assert measured["missing"].null_rate == 1.0 # all-null, kept for the null rate alone + assert set(probes) == {feature.name for feature in features} # every column, typed or not def test_a_column_with_nothing_to_measure_is_not_reported_as_an_error(): From f369b32f60860253f5e5243ea5b6da3e39497a70 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 18:59:55 -0400 Subject: [PATCH 38/44] feat(datasets): read quantiles off counters instead of retained lengths Phase 3, third step, and the reservoir the spec called for is not what shipped. Exact quantiles need every length kept and sorted, which is a list that grows with the dataset. A reservoir of sampled lengths bounds that, but buys the bound with an RNG -- and so with `SamplingInfo.seed` back in the contract after it was deleted on the grounds that the profiler makes no random choices. Counting into fixed buckets bounds it with neither. The two put their error in different places, and that is the whole argument. A reservoir sees *some* rows exactly: its error is in which rows it kept, which is probabilistic and only shrinks with the sample size. A histogram sees *every* row imprecisely: its error is in how finely each value was recorded, which is a hard bound of half a bucket width whatever the data does. Rounding the value is the cheap error to accept, because the number is read to pick a sequence budget and gets rounded to a power of two by whoever reads it. Lengths below 32 get a counter each and stay exact. Above that each octave is cut into 32 slices, so a bucket spans a fixed 1/32 of its value and the midpoint lands within ~1.6%. Measured against exact quantiles on the real shards, worst error over fifteen estimates is 1.5%, and the bucket/bounds round trip is verified from 0 to 33.5M. Midpoint, not the bucket's low edge. The edge sits systematically under the truth -- every estimate in the first measurement came out low -- and centring roughly halves the average error. Clamped to `max`, which is tracked separately and stays exact: a p99 above the largest value present would be nonsense, and `max` is the one number here a reader may treat as a hard bound. A messages column is now fully O(1): its two histograms cost 280 bytes and 22 KB over 43,835 rows, against ~342 KB for a retained list of lengths or ~800 KB for a 100k reservoir. A string column has one term left -- it still retains its strings, because the quality stride needs the column's row count to place its sample and that is not known until the last batch. Summing the parquet footers before folding removes it, and that is the next step. Contract: p50/p95/p99 are documented as estimates and `max` as exact. `p95` is kept -- I proposed dropping it and it was not part of what was agreed. `stats_complete` now says plainly that it speaks to rows read rather than to every number being exact, which replaces the narrower TextQuality-only carve-out from Phase 2. Also deleted `_message_stats`, superseded by MessageAccumulator two commits ago and still calling `_quantiles`. Tests never caught it because nothing called it; ruff and ty did. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 22 ++- .../nemo_datasets_plugin/profiler/stats.py | 170 ++++++++++-------- plugins/nemo-datasets/tests/test_stats.py | 45 +++++ 3 files changed, 161 insertions(+), 76 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index cfc42110c4..2b83760a4e 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -134,7 +134,21 @@ class PartitionClassification(BaseModel): class Quantiles(BaseModel): - """A per-row distribution summary. p99 = long-tail sequence-length signal; max = hard cap.""" + """A per-row distribution summary. p99 = long-tail sequence-length signal; max = hard cap. + + The shape is the point, not the precision. Mean and max cannot tell "uniformly medium-length" + apart from "mostly short with a long tail", and those call for opposite sequence budgets — set + one from `max` and most of the memory is wasted, set it from the mean and the tail is silently + truncated. Reading p50 against p99 is what answers it. + + **p50 / p95 / p99 are estimates, within a couple of percent.** They are read off counters bucketed + by magnitude rather than from the lengths themselves, which is what keeps the profiler's memory + flat in rows. Every row is counted, so the *rank* is exact; only the value is rounded, and it is + rounded to a bound that does not grow with the dataset. That is the cheap error to accept here, + because whoever reads these rounds to a power of two anyway. + + **`max` is exact**, always, and is the only number here safe to treat as a hard bound. + """ p50: int p95: int @@ -458,9 +472,9 @@ class PartitionProfile(BaseModel): "True => `features`, `stats` and `classification` were computed over every row of every " "file in THIS partition: proven facts, not estimates. Only then can a consumer assert " "enum / required in a bridged JSON Schema, or read a verifiability coverage of 1.0 as " - "literal. `TextQuality` is the one exception and is always an estimate — it is bounded " - "independently of how much was read, for cost reasons its own docstring gives — though " - "the two only diverge on an unbounded read of a column larger than that bound. " + "literal. It speaks to *rows read*, not to every number being exact: `TextQuality` and " + "`Quantiles` are estimates by construction however much was read, each bounded for the " + "cost reasons its own docstring gives, and `Quantiles.max` is exact regardless. " "Scoped to the partition because that is where it is decided — a corrupt shard " "in one partition says nothing about the measurements in another, and a fileset-wide " "flag quietly downgraded every partition to the worst one." diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 1b38166a0c..3e8f7bc62f 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -22,10 +22,10 @@ partition before it can measure it. The base class is the entire measurement for a dtype with no statistics of its own, because the probes run over every column whatever its type. -Two things here are still sized by the column rather than bounded: the retained strings a quality -stride needs to place its sample, and the lengths a quantile needs to sort. Both are what the -reservoir and the parquet footer row counts are for, and until then they cost what materialising the -column already cost. +Every measurement is now O(1) in rows except one: a string column retains its strings, because the +quality stride needs the column's row count to place its sample and that is not known until the last +batch. Summing the parquet footers before folding is what removes it — the row count is in them, and +reading it costs no rows. A messages column has no such term and is already bounded. """ from __future__ import annotations @@ -272,19 +272,25 @@ class StringAccumulator(ColumnAccumulator): def __init__(self) -> None: super().__init__() self._strings: list[str] = [] + self._lengths = _LengthHistogram() self._vocabulary = _Vocabulary() def _observe(self, present: list[Any]) -> None: - # The strings are retained because the quality stride needs the column's length to place its - # sample, and that is not known until the last batch. Bounding this is what the reservoir and - # the footer row count are for; until then it costs what materialising the column already did. - self._strings.extend(value for value in present if isinstance(value, str)) + for value in present: + if isinstance(value, str): + # The length folds away; the string itself is retained only because the quality + # stride needs the column's row count to place its sample, and that is not known + # until the last batch. Bounding that is the parquet footer sum's job -- it is the + # one term here still sized by the column, and it costs what materialising the + # column already cost. + self._strings.append(value) + self._lengths.add(len(value)) self._vocabulary.update(present) def _blocks(self) -> dict[str, Any]: text = quality = None if self._strings: - text = TextStats(chars=_quantiles([len(value) for value in self._strings])) + text = TextStats(chars=self._lengths.quantiles()) quality = _text_quality(self._strings) return {"text": text, "quality": quality, "categorical": self._vocabulary.finalize()} @@ -341,8 +347,8 @@ class MessageAccumulator(ColumnAccumulator): def __init__(self) -> None: super().__init__() self._conversations = 0 - self._turns: list[int] = [] - self._content_chars: list[int] = [] + self._turns = _LengthHistogram() + self._content_chars = _LengthHistogram() self._roles_seen: list[str] = [] self._ends_with_assistant = 0 self._valid_alternation = 0 @@ -353,7 +359,7 @@ def _observe(self, present: list[Any]) -> None: if not isinstance(messages, list): continue self._conversations += 1 - self._turns.append(len(messages)) + self._turns.add(len(messages)) total_content = 0 for message in messages: if not isinstance(message, dict): @@ -372,7 +378,7 @@ def _observe(self, present: list[Any]) -> None: # schema that merely declares tool_calls would otherwise report tool use on every row. if message.get("tool_calls") or role == "tool": self._has_tool_calls = True - self._content_chars.append(total_content) + self._content_chars.add(total_content) if messages and isinstance(messages[-1], dict) and _is_assistant_role(_role_of(messages[-1])): self._ends_with_assistant += 1 if _valid_alternation(messages): @@ -383,8 +389,8 @@ def _blocks(self) -> dict[str, Any]: return {"messages": None} return { "messages": MessageStats( - turns=_quantiles(self._turns), - content_chars=_quantiles(self._content_chars), + turns=self._turns.quantiles(), + content_chars=self._content_chars.quantiles(), roles_seen=self._roles_seen, ends_with_assistant_rate=self._ends_with_assistant / self._conversations, valid_alternation_rate=self._valid_alternation / self._conversations, @@ -410,18 +416,85 @@ def _is_numeric(dtype: str) -> bool: return dtype.startswith(("int", "uint", "float")) -def _quantiles(values: list[int]) -> Quantiles: - """Nearest-rank percentiles over the sample (n is small, so this stays exact).""" - ordered = sorted(values) - n = len(ordered) +# How finely a length distribution is recorded. Lengths below the slice count get a counter each and +# are exact; above it, each octave is cut into this many slices, so a bucket spans a fixed *relative* +# width of 1/32. Reporting a bucket's midpoint then puts every estimate within ~1.6% of the truth, +# whatever the value's magnitude and however many rows there are. +_HISTOGRAM_SLICE_BITS = 5 +_HISTOGRAM_SLICES = 1 << _HISTOGRAM_SLICE_BITS - def at(percentile: int) -> int: - if n == 0: - return 0 - rank = math.ceil(percentile / 100 * n) - return ordered[min(rank, n) - 1] - return Quantiles(p50=at(50), p95=at(95), p99=at(99), max=ordered[-1] if ordered else 0) +def _length_bucket(value: int) -> int: + """The counter a length belongs to.""" + if value < _HISTOGRAM_SLICES: + return value # small lengths get a counter each, so they are recorded exactly + shift = value.bit_length() - 1 - _HISTOGRAM_SLICE_BITS + return (shift + 1) * _HISTOGRAM_SLICES + ((value >> shift) - _HISTOGRAM_SLICES) + + +def _bucket_bounds(bucket: int) -> tuple[int, int]: + """The half-open range of lengths that land in ``bucket``. Inverse of :func:`_length_bucket`.""" + if bucket < _HISTOGRAM_SLICES: + return bucket, bucket + 1 + index, slice_index = divmod(bucket, _HISTOGRAM_SLICES) + shift = index - 1 + low = (_HISTOGRAM_SLICES + slice_index) << shift + return low, low + (1 << shift) + + +class _LengthHistogram: + """A per-row length distribution, held as counters rather than as the lengths themselves. + + This is what lets an accumulator stay O(1) in rows. Exact quantiles need every length kept and + sorted, which is a list that grows with the dataset; a reservoir of sampled lengths bounds that, + but buys the bound with an RNG -- and so with a seed back in the contract, and two runs over the + same bytes disagreeing. Counting into fixed buckets bounds it with neither. + + The two put their error in different places. A reservoir sees *some* rows exactly, so its error + is in which rows it happened to keep: probabilistic, and shrinking only with the sample size. + This sees *every* row imprecisely, so its error is in how finely each value was recorded: a hard + bound of half a bucket width, whatever the data does. Measured against exact quantiles on real + shards, ~2%. + + Rounding the value is the cheap error to accept here, because the number is read to pick a + sequence budget and gets rounded to a power of two by whoever reads it. ``max`` is kept exactly + and separately: it is the one value here a reader may treat as a hard bound. + """ + + def __init__(self) -> None: + self._counts: dict[int, int] = {} + self._rows = 0 + self._max = 0 + + def add(self, value: int) -> None: + bucket = _length_bucket(value) + self._counts[bucket] = self._counts.get(bucket, 0) + 1 + self._rows += 1 + if value > self._max: + self._max = value + + def quantiles(self) -> Quantiles: + return Quantiles(p50=self._at(50), p95=self._at(95), p99=self._at(99), max=self._max) + + def _at(self, percentile: int) -> int: + """Nearest-rank percentile: the bucket the p-th row falls in, reported at its midpoint. + + The rank is exact -- every row was counted, none sampled -- so only the value is approximate. + Midpoint rather than the bucket's low edge, which sits systematically under the truth and + roughly doubles the average error. + """ + if not self._rows: + return 0 + target = math.ceil(percentile / 100 * self._rows) + seen = 0 + for bucket in sorted(self._counts): + seen += self._counts[bucket] + if seen >= target: + low, high = _bucket_bounds(bucket) + # Never above `max`: a midpoint can overshoot the largest value actually present, + # and a p99 above the maximum would be nonsense. + return min((low + high) // 2, self._max) + return self._max # --- text quality -------------------------------------------------------------------------------- @@ -560,53 +633,6 @@ def _is_assistant_role(role: Any) -> bool: return isinstance(role, str) and role.lower() in _ASSISTANT_ROLES -def _message_stats(rows_messages: list[list]) -> MessageStats | None: - if not rows_messages: - return None - turns: list[int] = [] - content_chars: list[int] = [] - roles_seen: list[str] = [] - ends_with_assistant = 0 - valid_alternation = 0 - has_tool_calls = False - - for messages in rows_messages: - turns.append(len(messages)) - total_content = 0 - for message in messages: - if not isinstance(message, dict): - continue - role = _role_of(message) - if role is not None: - # Coerced to str because roles_seen is typed list[str] and a non-string role would - # fail validation — aborting the whole profile from inside the one stage the pipeline - # does not guard. Reported verbatim otherwise: the contract is explicit that an - # unexpected role is the finding worth surfacing, not something to normalize away. - role = role if isinstance(role, str) else str(role) - if role not in roles_seen: - roles_seen.append(role) - total_content += _content_len(_message_field(message, "content", "value")) - # `.get` truthiness, not `in`: parquet materializes every declared struct field, so a - # schema that merely declares tool_calls would otherwise report tool use on every row. - if message.get("tool_calls") or role == "tool": - has_tool_calls = True - content_chars.append(total_content) - if messages and isinstance(messages[-1], dict) and _is_assistant_role(_role_of(messages[-1])): - ends_with_assistant += 1 - if _valid_alternation(messages): - valid_alternation += 1 - - n = len(rows_messages) - return MessageStats( - turns=_quantiles(turns), - content_chars=_quantiles(content_chars), - roles_seen=roles_seen, - ends_with_assistant_rate=ends_with_assistant / n, - valid_alternation_rate=valid_alternation / n, - has_tool_calls=has_tool_calls, - ) - - def _content_len(content: Any) -> int: if isinstance(content, str): return len(content) diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index ec8cc60c62..11fe8ad894 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -3,6 +3,8 @@ """Tests for per-column statistics.""" +import math + import pytest from nemo_datasets_plugin.profiler import stats as stats_module from nemo_datasets_plugin.profiler.stats import ( @@ -38,6 +40,49 @@ def _rows(name, values): return [{name: value} for value in values] +# --- length histogram ---------------------------------------------------------------------------- + + +@pytest.mark.parametrize("value", [0, 1, 31, 32, 33, 63, 64, 255, 256, 1_000, 1_300, 65_535, 1_000_000, 33_554_432]) +def test_every_length_lands_in_a_bucket_that_contains_it(value): + # The bounds are what a quantile is read off, so they have to invert the bucketing exactly. A + # bucket whose range does not contain its own values would report a plausible wrong number. + low, high = stats_module._bucket_bounds(stats_module._length_bucket(value)) + assert low <= value < high + + +def test_short_lengths_are_recorded_exactly(): + # Below the slice count every length gets its own counter. That is what keeps the small fixtures + # in this file exact, and it is why a column of short strings loses nothing to bucketing. + hist = stats_module._LengthHistogram() + for n in range(stats_module._HISTOGRAM_SLICES): + hist.add(n) + quantiles = hist.quantiles() + assert (quantiles.p50, quantiles.p95, quantiles.p99, quantiles.max) == (15, 30, 31, 31) + + +def test_quantiles_stay_within_the_bound_on_a_heavy_tail(): + # The shape that matters: most rows short, a thin long tail. It is also the shape a mean cannot + # describe, which is why the distribution is carried at all. + values = [10] * 5000 + [200] * 3000 + [4000] * 1500 + [90_000] * 500 + hist = stats_module._LengthHistogram() + for value in values: + hist.add(value) + quantiles = hist.quantiles() + + ordered = sorted(values) + for percentile, got in ((50, quantiles.p50), (95, quantiles.p95), (99, quantiles.p99)): + want = ordered[min(math.ceil(percentile / 100 * len(ordered)), len(ordered)) - 1] + assert abs(got - want) / want <= 0.02, (percentile, want, got) + assert quantiles.max == 90_000 # exact, never rounded to a bucket + assert quantiles.p50 <= quantiles.p95 <= quantiles.p99 <= quantiles.max + + +def test_an_empty_histogram_reports_zeros(): + quantiles = stats_module._LengthHistogram().quantiles() + assert (quantiles.p50, quantiles.p95, quantiles.p99, quantiles.max) == (0, 0, 0, 0) + + # --- text ---------------------------------------------------------------------------------------- From 1c978750c54bcc08de848823ca6b17c40daf8944 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 19:12:42 -0400 Subject: [PATCH 39/44] refactor(datasets): take rows out of the measure stage Phase 3, fourth step. `quote_enumerations` and `classify` were the last two things between a measured partition and a materialised one, and both were reaching back to the rows for something already computed. `quote_enumerations` rescanned every row to rebuild a distinct set the vocabulary had just built and thrown away. `_Vocabulary` now hands out what it kept and the pass reads that. `measure_columns` grew a named result to carry it -- three positional returns was already the limit, and the vocabularies are not part of the stored profile, so a tuple was the wrong shape for them. `classify` took rows for two reasons. It derived probes when handed none, which only the tests ever used; absent probes now read as "nothing was measured", which is the honest reading and never "nothing is there". And it ran the chosen/rejected shared-prefix check, which is the one probe that compares two columns of the same row against each other and so can live neither on a column accumulator nor in a per-column probe. That becomes `PrefixPairFold`: two counters, folded over the same batches, resolved by column *name* because the fold runs before classification has assigned any roles -- the same inversion the content probes made when they stopped being role-gated. `derive_probes` goes with it, having lost its only production caller for the same reason `derive_stats` did two commits ago. The tests keep a two-line local view. The prefix threshold, inlined as `16`, is now named. Nothing in the measure stage reads a row any more except the two folds themselves, which is what the next step needs: batches arrive, both folds consume them, and schema, stats, probes, classification and quoting all follow from what was folded. Behaviour is unchanged and checked as such: the demo profile is byte-identical to the one HEAD produces, diffed against a stashed build of the same tree. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/classify.py | 97 +++++++++++++------ .../nemo_datasets_plugin/profiler/pipeline.py | 22 +++-- .../nemo_datasets_plugin/profiler/stats.py | 87 +++++++++++------ plugins/nemo-datasets/tests/test_classify.py | 56 +++++++---- plugins/nemo-datasets/tests/test_stats.py | 41 ++++---- 5 files changed, 198 insertions(+), 105 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index 9426531c20..311a5061c4 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -15,7 +15,9 @@ from __future__ import annotations -from nemo_datasets_plugin.profiler.stats import ColumnProbes, derive_probes +from dataclasses import dataclass + +from nemo_datasets_plugin.profiler.stats import ColumnProbes from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, Evidence, @@ -317,6 +319,12 @@ def _detect_verifiability(features: list[FeatureSchema], probes: dict[str, Colum return None +# How much text two answers must open with in common before it reads as a shared prompt rather than +# a shared turn of phrase. Short enough that a one-line question counts, long enough that "I think +# that" does not. +_EMBEDDED_PROMPT_PREFIX_CHARS = 16 + + def _common_prefix_len(left: str, right: str) -> int: limit = min(len(left), len(right)) index = 0 @@ -325,8 +333,53 @@ def _common_prefix_len(left: str, right: str) -> int: return index +@dataclass(frozen=True) +class PrefixPair: + """How often two columns of the same row opened with the same long run of text.""" + + pairs: int = 0 + shared: int = 0 + + +class PrefixPairFold: + """The one probe that reads two columns against each other rather than each on its own. + + A preference set whose prompt is embedded in both answers shows up as a long shared prefix + between them, and no per-column measurement can see that. Being relational, it also cannot live + on a column accumulator, so it folds separately over the same batches. + + Resolved by column *name*, not by role, because the fold runs before classification has assigned + any roles -- the same inversion the content probes made when they stopped being role-gated. The + names are the ones the alias table maps to these two roles. + """ + + def __init__(self, features: list[FeatureSchema]) -> None: + self._left = self._column_named(features, "chosen") + self._right = self._column_named(features, "rejected") + self._pairs = 0 + self._shared = 0 + + @staticmethod + def _column_named(features: list[FeatureSchema], role: str) -> str | None: + wanted = {name for name, aliased in _ALIAS_ROLES.items() if aliased == role} + return next((f.name for f in features if f.name in wanted and f.dtype == "string"), None) + + def update(self, rows: list[dict]) -> None: + if self._left is None or self._right is None: + return + for row in rows: + left, right = row.get(self._left), row.get(self._right) + if isinstance(left, str) and isinstance(right, str): + self._pairs += 1 + if _common_prefix_len(left, right) >= _EMBEDDED_PROMPT_PREFIX_CHARS: + self._shared += 1 + + def result(self) -> PrefixPair: + return PrefixPair(pairs=self._pairs, shared=self._shared) + + def _implicit_prompt_evidence( - features: list[FeatureSchema], probes: dict[str, ColumnProbes], rows: list[dict] + features: list[FeatureSchema], probes: dict[str, ColumnProbes], prefix_pair: PrefixPair ) -> Evidence | None: targets = [f for f in features if f.semantic_role in {"chosen", "rejected", "completion"} and f.dtype == "string"] counted = [probes[f.name] for f in targets if f.name in probes] @@ -336,49 +389,35 @@ def _implicit_prompt_evidence( detail = f"embedded transcript markers in {_pct(marked / sampled)} of sampled completions - prompt is embedded" return Evidence(kind="content_probe", detail=detail) - # The shared-prefix check is *relational* — it compares two columns against each other — so it - # has no per-column probe to read and still works from the rows themselves. - if not rows: - return None - - chosen = next((f for f in features if f.semantic_role == "chosen" and f.dtype == "string"), None) - rejected = next((f for f in features if f.semantic_role == "rejected" and f.dtype == "string"), None) - if chosen is not None and rejected is not None: - pairs = 0 - shared = 0 - for row in rows: - left, right = row.get(chosen.name), row.get(rejected.name) - if isinstance(left, str) and isinstance(right, str): - pairs += 1 - if _common_prefix_len(left, right) >= 16: - shared += 1 - if pairs and shared / pairs >= 0.5: - detail = f"chosen/rejected share a common prefix in {_pct(shared / pairs)} of pairs - prompt is embedded" - return Evidence(kind="content_probe", detail=detail) + if prefix_pair.pairs and prefix_pair.shared / prefix_pair.pairs >= 0.5: + rate = _pct(prefix_pair.shared / prefix_pair.pairs) + return Evidence( + kind="content_probe", + detail=f"chosen/rejected share a common prefix in {rate} of pairs - prompt is embedded", + ) return None def classify( features: list[FeatureSchema], stats: dict[str, ColumnStats], - rows: list[dict] | None = None, *, probes: dict[str, ColumnProbes] | None = None, + prefix_pair: PrefixPair | None = None, column_roles: dict[str, str] | None = None, ) -> PartitionClassification: """Assign roles onto ``features`` in place and return the partition's classification. - ``probes`` are the per-column content measurements from :func:`~.stats.derive_probes`. They are - a pure function of ``(features, rows)``, so a caller that has not already computed them can pass - ``rows`` alone and get them derived here; the pipeline passes them in to avoid the second pass. - Role/axis/type inference needs neither — only the schema and stats. + ``probes`` are the per-column content measurements, and ``prefix_pair`` the one relational one. + Both are folded over the rows before this runs; nothing here reads a row, which is what lets a + partition be classified without ever having been materialised. Absent, each reads as "nothing + was measured", and role/axis/type inference is unaffected — that needs only schema and stats. ``column_roles`` maps a column name to a role the caller is asserting, taking precedence over the name-alias table but still subject to the dtype gates. It exists because that table is ~35 English names with no way to say "my `q` column is the prompt", and its misses are silent. """ - rows = rows or [] - probes = derive_probes(features, rows) if probes is None else probes + probes = probes or {} evidence = _assign_roles(features, stats, column_roles or {}) roles = {feature.semantic_role for feature in features if feature.semantic_role} candidates = _detect_types(features, stats) @@ -392,7 +431,7 @@ def classify( if fmt is not None: evidence.append(Evidence(kind="column_dtype", detail=f"{fmt} format from role column dtypes")) if prompt_form == "implicit": - embedded = _implicit_prompt_evidence(features, probes, rows) + embedded = _implicit_prompt_evidence(features, probes, prefix_pair or PrefixPair()) if embedded is not None: evidence.append(embedded) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index b736ff52c0..ff2428d25c 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -29,7 +29,7 @@ from pathlib import PurePosixPath import pyarrow as pa -from nemo_datasets_plugin.profiler.classify import classify +from nemo_datasets_plugin.profiler.classify import PrefixPairFold, classify from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource from nemo_datasets_plugin.profiler.partition import group_partitions from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data @@ -244,15 +244,25 @@ def _measure( # Statistics and probes together, one column at a time and each isolated. Probes cover every # column independent of the roles classify is about to assign, so a content signal survives # a column name the alias table does not know. - stats, probes, column_errors = measure_columns(features, partition_rows) - classification = classify(features, stats, partition_rows, probes=probes, column_roles=column_roles) + measured = measure_columns(features, partition_rows) + # The one probe that reads two columns against each other, so it cannot live on a column's + # accumulator. Folded here over the same rows. + prefix_pair = PrefixPairFold(features) + prefix_pair.update(partition_rows) + classification = classify( + features, + measured.stats, + probes=measured.probes, + prefix_pair=prefix_pair.result(), + column_roles=column_roles, + ) # Last, because the roles classification assigns are what decide whether a column's values # may be quoted at all — cardinality only bounds how many. - quote_enumerations(features, stats, partition_rows) + quote_enumerations(features, measured.stats, measured.vocabularies) # After classify, so its own reasoning reads first and a column that could not be measured # is reported as a caveat on the result rather than as part of the case for it. - classification.evidence.extend(column_errors) - return features, stats, classification + classification.evidence.extend(measured.errors) + return features, measured.stats, classification except Exception as exc: detail = f"could not measure this partition: {type(exc).__name__}: {exc}" return [], {}, PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 3e8f7bc62f..f9d0010a67 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -11,10 +11,9 @@ small controlled vocabulary is added afterwards by :func:`quote_enumerations`, which gates on role. The same pass reads each column's *content* — answer markers, embedded transcripts — as plain -per-column counts (:class:`ColumnProbes`, also reachable alone via :func:`derive_probes`, which -classification uses when it was handed no probes). Those are measurements, not interpretations: what -they mean is classification's job, and keeping the looking here is what stops a content signal from -being reachable only through a correctly named column. +per-column counts (:class:`ColumnProbes`). Those are measurements, not interpretations: what they +mean is classification's job, and keeping the looking here is what stops a content signal from being +reachable only through a correctly named column. The measuring itself is done by a :class:`ColumnAccumulator` per column, chosen once on dtype. An accumulator folds batches in and keeps no reference to them, so a column measured in pieces gives @@ -71,9 +70,22 @@ _QUOTABLE_ROLES = frozenset({"label", "provenance", "meta", "rank"}) -def measure_columns( - features: list[FeatureSchema], rows: list[dict[str, Any]] -) -> tuple[dict[str, ColumnStats], dict[str, ColumnProbes], list[Evidence]]: +@dataclass(frozen=True) +class ColumnMeasurements: + """Everything one pass over a partition's columns produced. + + A named result rather than a tuple because the vocabularies joined it: they are not part of the + stored profile, but :func:`quote_enumerations` needs them and can no longer go back to the rows + for them. + """ + + stats: dict[str, ColumnStats] + probes: dict[str, ColumnProbes] + vocabularies: dict[str, set[Any]] # name -> distinct values, only where the column stayed one + errors: list[Evidence] + + +def measure_columns(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> ColumnMeasurements: """Measure every top-level column: its statistics and its content probes, in one pass each. Each column is isolated. A value no detector anticipated -- a chat message whose ``role`` is a @@ -93,6 +105,7 @@ def measure_columns( """ stats: dict[str, ColumnStats] = {} probes: dict[str, ColumnProbes] = {} + vocabularies: dict[str, set[Any]] = {} errors: list[Evidence] = [] for feature in features: # Parquet permits duplicate field names, and both maps are keyed by name. Measuring the @@ -108,6 +121,9 @@ def measure_columns( accumulator.update([row.get(feature.name) for row in rows]) column, probe = accumulator.finalize() probes[feature.name] = probe + vocabulary = accumulator.vocabulary() + if vocabulary is not None: + vocabularies[feature.name] = vocabulary except Exception as exc: errors.append( Evidence( @@ -120,11 +136,11 @@ def measure_columns( continue if column is not None: stats[feature.name] = column - return stats, probes, errors + return ColumnMeasurements(stats=stats, probes=probes, vocabularies=vocabularies, errors=errors) def quote_enumerations( - features: list[FeatureSchema], stats: dict[str, ColumnStats], rows: list[dict[str, Any]] + features: list[FeatureSchema], stats: dict[str, ColumnStats], vocabularies: dict[str, set[Any]] ) -> None: """Fill in ``categorical.values`` for columns whose role makes them a controlled vocabulary. @@ -132,6 +148,10 @@ def quote_enumerations( the way classification mutates ``features``. Deliberately fills in rather than redacting: skip this pass and no values are stored, where a redaction pass that got skipped would leak them. + Reads what the accumulators already kept rather than going back to the rows. That second pass + was the last thing tying the measure stage to a materialised partition, and it was re-deriving a + set the vocabulary had built and thrown away. + Cardinality is only the size bound. It cannot be the permission, because it inverts on small data -- in a three-row dataset every column holds under 32 distinct values, free text included, so an entire column of prompts was quotable. The role says what a column *is*, at any size. @@ -142,11 +162,10 @@ def quote_enumerations( column = stats.get(feature.name) if column is None or column.categorical is None or column.categorical.distinct_count > _MAX_ENUM_VALUES: continue - try: - distinct = {value for row in rows if (value := row.get(feature.name)) is not None} - except TypeError: - continue # unhashable values have no enumeration to quote - column.categorical.values = sorted(str(value) for value in distinct) + values = vocabularies.get(feature.name) + if values is None: + continue + column.categorical.values = sorted(str(value) for value in values) class ColumnAccumulator: @@ -215,6 +234,11 @@ def _blocks(self) -> dict[str, Any]: """The dtype-specific ``ColumnStats`` blocks. The base column contributes none.""" return {} + def vocabulary(self) -> set[Any] | None: + """The distinct values, for a column that is a bounded vocabulary. None for one that is not, + which is every dtype without a notion of cardinality.""" + return None + class _Vocabulary: """Distinct values, for as long as the column still looks like a controlled vocabulary. @@ -265,6 +289,14 @@ def _saturate(self) -> None: def finalize(self) -> CategoricalStats | None: return None if self._saturated else CategoricalStats(distinct_count=len(self._values)) + def values(self) -> set[Any] | None: + """What was kept, or None once the column stopped being a vocabulary. + + Handing this out is what lets :func:`quote_enumerations` fill in an enumeration without a + second pass over the rows -- which it could only do while the rows were still there. + """ + return None if self._saturated else self._values + class StringAccumulator(ColumnAccumulator): """A ``string`` column: length quantiles, corruption ratios, and a vocabulary if it has one.""" @@ -294,6 +326,9 @@ def _blocks(self) -> dict[str, Any]: quality = _text_quality(self._strings) return {"text": text, "quality": quality, "categorical": self._vocabulary.finalize()} + def vocabulary(self) -> set[Any] | None: + return self._vocabulary.values() + class NumericAccumulator(ColumnAccumulator): """An ``int*`` / ``uint*`` / ``float*`` column: running extrema and mean, plus a vocabulary.""" @@ -325,6 +360,9 @@ def _blocks(self) -> dict[str, Any]: numeric = NumericStats(min=self._min, max=self._max, mean=self._sum / self._count) return {"numeric": numeric, "categorical": self._vocabulary.finalize()} + def vocabulary(self) -> set[Any] | None: + return self._vocabulary.values() + class BoolAccumulator(ColumnAccumulator): """A ``bool`` column. The column that decides unpaired_preference deserves a measured class @@ -340,6 +378,9 @@ def _observe(self, present: list[Any]) -> None: def _blocks(self) -> dict[str, Any]: return {"categorical": self._vocabulary.finalize()} + def vocabulary(self) -> set[Any] | None: + return self._vocabulary.values() + class MessageAccumulator(ColumnAccumulator): """A ``messages`` column: turn and length distributions, the roles seen, and chat-shape rates.""" @@ -678,24 +719,6 @@ class ColumnProbes: transcript_marker: int # of `texts`, how many embed a Human:/Assistant: transcript -def derive_probes(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> dict[str, ColumnProbes]: - """Run the content probes over every top-level column, keyed by column name. - - The probes alone, through the base accumulator, which carries no dtype state — so this costs the - scan and nothing else. Classification calls it when it was handed no probes of its own; - :func:`measure_columns` is the path that wants the statistics as well. - """ - probes: dict[str, ColumnProbes] = {} - for feature in features: - # Duplicate parquet field names: first wins, matching `measure_columns` so both agree on which. - if feature.name in probes: - continue - accumulator = ColumnAccumulator() - accumulator.update([row.get(feature.name) for row in rows]) - probes[feature.name] = accumulator.finalize()[1] - return probes - - def _probe_text(value: Any) -> str | None: """The text a probe reads from one cell: the string itself, or a chat column's final turn. diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index 05858e7438..ef1e72a533 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -3,8 +3,8 @@ """Tests for classification: role assignment, format/prompt-form axes, and dataset type.""" -from nemo_datasets_plugin.profiler.classify import classify -from nemo_datasets_plugin.profiler.stats import derive_probes +from nemo_datasets_plugin.profiler.classify import PrefixPairFold, classify +from nemo_datasets_plugin.profiler.stats import measure_columns from nemo_platform_plugin.files.dataset_profile import ( CategoricalStats, ColumnStats, @@ -14,6 +14,17 @@ ) +def _probes(features, rows): + return measure_columns(features, rows).probes + + +def classify_rows(features, stats, rows, **kwargs): + """Classify from rows the way the pipeline does: fold first, then interpret the folds.""" + prefix = PrefixPairFold(features) + prefix.update(rows) + return classify(features, stats, probes=_probes(features, rows), prefix_pair=prefix.result(), **kwargs) + + def _f(name, dtype): return FeatureSchema(name=name, dtype=dtype) @@ -174,14 +185,14 @@ def test_classification_records_evidence(): def test_verifiability_extractable_gsm8k_answer(): features = [_f("problem", "string"), _f("solution", "string")] rows = [{"problem": "q", "solution": "steps #### 18"}, {"problem": "q", "solution": "no final answer"}] - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert result.verifiability.method == "extractable_final_answer" assert result.verifiability.coverage == 0.5 def test_verifiability_boxed_answer(): features = [_f("prompt", "string"), _f("completion", "string")] - result = classify(features, {}, [{"prompt": "q", "completion": r"reasoning \boxed{42}"}]) + result = classify_rows(features, {}, [{"prompt": "q", "completion": r"reasoning \boxed{42}"}]) assert result.verifiability.method == "extractable_final_answer" assert result.verifiability.coverage == 1.0 @@ -189,14 +200,14 @@ def test_verifiability_boxed_answer(): def test_verifiability_ground_truth_column_coverage(): features = [_f("prompt", "string"), _f("ground_truth", "string")] rows = [{"prompt": "q", "ground_truth": "42"}, {"prompt": "q", "ground_truth": None}] - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert result.verifiability.method == "ground_truth_column" assert result.verifiability.coverage == 0.5 def test_no_verifiability_without_a_target(): features = [_f("prompt", "string"), _f("completion", "string")] - result = classify(features, {}, [{"prompt": "q", "completion": "just prose, no answer"}]) + result = classify_rows(features, {}, [{"prompt": "q", "completion": "just prose, no answer"}]) assert result.verifiability is None @@ -205,7 +216,7 @@ def test_verifiability_ignores_below_threshold_extractable_noise(): features = [_f("prompt", "string"), _f("completion", "string")] rows = [{"prompt": "q", "completion": "just prose"} for _ in range(100)] rows[0]["completion"] = "the answer is #### 7" # 1/100 = 1% < 5% floor - assert classify(features, {}, rows).verifiability is None + assert classify_rows(features, {}, rows).verifiability is None def test_verifiability_asserted_above_coverage_floor(): @@ -213,7 +224,7 @@ def test_verifiability_asserted_above_coverage_floor(): rows = [{"prompt": "q", "completion": "just prose"} for _ in range(10)] for row in rows[:2]: row["completion"] = "answer #### 7" # 2/10 = 20% >= 5% floor - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert result.verifiability.method == "extractable_final_answer" assert result.verifiability.coverage == 0.2 @@ -223,7 +234,7 @@ def test_sparse_ground_truth_falls_through_to_extractable_answer(): features = [_f("completion", "string"), _f("ground_truth", "string")] rows = [{"completion": "reasoning #### 5", "ground_truth": None} for _ in range(100)] rows[0]["ground_truth"] = "5" # 1/100 ground_truth coverage -> below floor, must fall through - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert result.verifiability.method == "extractable_final_answer" assert result.verifiability.coverage == 1.0 @@ -231,7 +242,7 @@ def test_sparse_ground_truth_falls_through_to_extractable_answer(): def test_implicit_prompt_evidence_from_embedded_transcript(): features = [_f("chosen", "string"), _f("rejected", "string")] rows = [{"chosen": "\n\nHuman: hi\n\nAssistant: hello", "rejected": "\n\nHuman: hi\n\nAssistant: hey"}] - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert result.prompt_form == "implicit" assert any(e.kind == "content_probe" for e in result.evidence) @@ -248,7 +259,7 @@ def test_ground_truth_may_be_a_container_dtype(): def test_container_ground_truth_drives_verifiability(): features = [_f("prompt", "string"), _f("test_cases", "list")] rows = [{"prompt": "q", "test_cases": [{"in": "1", "out": "2"}]}, {"prompt": "q2", "test_cases": []}] - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert result.verifiability.method == "ground_truth_column" assert result.verifiability.coverage == 0.5 # the empty test_cases list is not a usable target @@ -266,7 +277,7 @@ def test_verifiability_survives_an_unrecognized_column_name(): # nothing knew where to look. The finding must name the column it came from. features = [_f("q", "string"), _f("a", "string")] rows = [{"q": "what is 2+2?", "a": f"add them #### {i}"} for i in range(10)] - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert {f.semantic_role for f in features} == {None} # still unroled, and honest about it assert result.dataset_type == "unknown" @@ -280,26 +291,29 @@ def test_a_named_completion_still_decides_where_to_look(): # completion is a better answer than one that merely looks like it. features = [_f("completion", "string"), _f("notes", "string")] rows = [{"completion": "just prose", "notes": "scratch #### 9"} for _ in range(10)] - assert classify(features, {}, rows).verifiability is None + assert classify_rows(features, {}, rows).verifiability is None def test_verifiability_reads_a_sharegpt_conversational_completion(): features = [_f("prompt", "string"), _f("completion", "messages")] rows = [{"prompt": "q", "completion": [{"from": "human", "value": "q"}, {"from": "gpt", "value": "#### 4"}]}] - result = classify(features, {}, rows) + result = classify_rows(features, {}, rows) assert result.verifiability.method == "extractable_final_answer" assert result.verifiability.coverage == 1.0 -def test_precomputed_probes_and_derived_probes_agree(): - # classify() derives probes when it is not given them; the pipeline passes them in to avoid the - # second pass. The two paths must not drift. +def test_classification_without_probes_claims_nothing_rather_than_guessing(): + # classify() used to derive probes from rows when it was handed none. It no longer sees rows at + # all, so absent probes must read as "nothing was measured" -- never as "nothing is there". rows = [{"prompt": "q", "completion": f"steps #### {i}"} for i in range(10)] - derived = classify([_f("prompt", "string"), _f("completion", "string")], {}, rows) features = [_f("prompt", "string"), _f("completion", "string")] - passed_in = classify(features, {}, rows, probes=derive_probes(features, rows)) - assert derived.verifiability.coverage == passed_in.verifiability.coverage - assert derived.verifiability.evidence[0].detail == passed_in.verifiability.evidence[0].detail + + blind = classify(features, {}) + assert blind.verifiability is None + + measured = classify([_f("prompt", "string"), _f("completion", "string")], {}, probes=_probes(features, rows)) + assert measured.verifiability.method == "extractable_final_answer" + assert measured.verifiability.coverage == 1.0 # --- candidates ------------------------------------------------------------------------------------ diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 11fe8ad894..e6f4e66b53 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -17,7 +17,6 @@ _non_ascii_count, _quality_sample, _whitespace_count, - derive_probes, measure_columns, quote_enumerations, ) @@ -27,9 +26,14 @@ def _stats(features, rows): """Statistics only. Asserts nothing failed: these tests measure values, not the guard, and a swallowed exception would surface here as a confusing KeyError instead of its own message.""" - measured, _, errors = measure_columns(features, rows) - assert not errors, errors - return measured + measured = measure_columns(features, rows) + assert not measured.errors, measured.errors + return measured.stats + + +def _probes(features, rows): + """The content probes alone. `measure_columns` measures both in one pass; these tests want one.""" + return measure_columns(features, rows).probes def _feature(name, dtype): @@ -117,7 +121,8 @@ def _observe(self, present): ) features = [_feature("good", "string"), _feature("bad", "string")] - measured, probes, errors = measure_columns(features, [{"good": "x", "bad": "y"}]) + result = measure_columns(features, [{"good": "x", "bad": "y"}]) + measured, probes, errors = result.stats, result.probes, result.errors assert "good" in measured and "bad" not in measured assert "good" in probes and "bad" not in probes # probes go with the column that failed @@ -169,10 +174,10 @@ def test_the_typed_accumulators_probe_exactly_as_the_bare_one_does(): features = [_feature(dtype, dtype) for dtype in sorted(_DTYPE_VALUES)] rows = [dict(zip(sorted(_DTYPE_VALUES), values)) for values in zip(*_DTYPE_VALUES.values())] - _, probes, errors = measure_columns(features, rows) + probes, errors = measure_columns(features, rows).probes, measure_columns(features, rows).errors assert errors == [] - assert probes == derive_probes(features, rows) # derive_probes goes through the base class + assert probes == _probes(features, rows) # probes come off the base class def test_measure_columns_measures_every_dtype_the_dispatch_knows(): @@ -197,7 +202,8 @@ def test_measure_columns_measures_every_dtype_the_dispatch_knows(): } for i in range(5) ] - measured, probes, errors = measure_columns(features, rows) + result = measure_columns(features, rows) + measured, probes, errors = result.stats, result.probes, result.errors assert errors == [] assert measured["text"].text is not None and measured["text"].quality is not None @@ -214,7 +220,8 @@ def test_measure_columns_measures_every_dtype_the_dispatch_knows(): def test_a_column_with_nothing_to_measure_is_not_reported_as_an_error(): # Absence from `stats` is the normal sparse case. Only a *failure* earns an error, or the two # would be indistinguishable and the guard would cry wolf on every well-formed struct column. - measured, probes, errors = measure_columns([_feature("s", "struct")], [{"s": {"a": 1}}]) + result = measure_columns([_feature("s", "struct")], [{"s": {"a": 1}}]) + measured, probes, errors = result.stats, result.probes, result.errors assert measured == {} and errors == [] assert "s" in probes @@ -437,7 +444,7 @@ def test_probes_are_measured_for_every_column_not_just_named_ones(): # alias table does not know still gets its content read. features = [_feature("q", "string"), _feature("a", "string")] rows = [{"q": "what is 2+2?", "a": "add them #### 4"}, {"q": "and 3+3?", "a": "no final answer"}] - probes = derive_probes(features, rows) + probes = _probes(features, rows) assert set(probes) == {"q", "a"} assert probes["a"].texts == 2 @@ -447,7 +454,7 @@ def test_probes_are_measured_for_every_column_not_just_named_ones(): def test_probes_read_the_final_turn_of_a_chat_column(): rows = [{"m": [{"role": "user", "content": "q"}, {"role": "assistant", "content": "steps #### 7"}]}] - probes = derive_probes([_feature("m", "messages")], rows) + probes = _probes([_feature("m", "messages")], rows) assert probes["m"].texts == 1 assert probes["m"].extractable_answer == 1 @@ -456,7 +463,7 @@ def test_probes_read_the_sharegpt_message_spelling(): # {from, value} is handled in schema derivation and message stats; reading only {role, content} # here cost every ShareGPT-shaped dataset its verifiability. rows = [{"m": [{"from": "human", "value": "q"}, {"from": "gpt", "value": "steps #### 7"}]}] - probes = derive_probes([_feature("m", "messages")], rows) + probes = _probes([_feature("m", "messages")], rows) assert probes["m"].texts == 1 assert probes["m"].extractable_answer == 1 @@ -466,14 +473,14 @@ def test_probes_count_non_empty_across_container_dtypes(): # target is just as often a list or struct as a string. features = [_feature("gt", "list")] rows = [{"gt": [{"in": "1"}]}, {"gt": []}, {"gt": None}] - probes = derive_probes(features, rows) + probes = _probes(features, rows) assert probes["gt"].rows == 3 assert probes["gt"].non_empty == 1 def test_probes_detect_embedded_transcripts(): rows = [{"c": "\n\nHuman: hi\n\nAssistant: hello"}, {"c": "plain prose"}] - probes = derive_probes([_feature("c", "string")], rows) + probes = _probes([_feature("c", "string")], rows) assert probes["c"].transcript_marker == 1 @@ -485,9 +492,9 @@ def _quoted(name, dtype, values, role): feature = _feature(name, dtype) feature.semantic_role = role rows = _rows(name, values) - stats = _stats([feature], rows) - quote_enumerations([feature], stats, rows) - return stats[name].categorical.values + measured = measure_columns([feature], rows) + quote_enumerations([feature], measured.stats, measured.vocabularies) + return measured.stats[name].categorical.values def test_quotes_a_controlled_vocabulary_role(): From a375d6dcad4746c296887566c1081fa07a86b8a5 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 19:24:28 -0400 Subject: [PATCH 40/44] feat(datasets): fold parquet partitions instead of materialising them Phase 3, last step. A partition whose files all declare a schema is now measured batch by batch and never held. The footers make it possible. A parquet file declares its schema and its exact row count there, so one seek per file establishes the partition's whole shape before a row is parsed -- which is precisely what a fold cannot otherwise have. The accumulators have to exist before the first batch, and the quality stride has to be placed before the column it strides has been seen. Both come out of `peek()`. Result on a real shard: 65.1 MB to profile 21,362 rows exhaustively becomes 10.4 MB, and 19.9 MB to profile 6,038 becomes 10.6 MB. Reading everything now costs what reading some of it costs. `row_budget` stops being a memory guard and becomes a limit on work. A string accumulator no longer retains its strings. With the row count known it places the quality stride up front and measures or skips each string as it goes by, holding counters instead. That was the last term sized by the column. Partitions without a declared schema still materialise, because the schema has to be inferred from the rows and the rows have to be kept until it has been. That is line-delimited formats; folding them needs accumulators created lazily as columns appear, which is Phase 4. Both readers grew `peek()` and `batches()` regardless, and the jsonl parse loop is now shared between `read()` and `batches()` so the two cannot drift on what counts as a row -- a blank line, a stray scalar and a truncated line are three different things and only one is an error. Behaviour is unchanged, checked two ways. The demo profile is byte-identical to HEAD's, diffed against a stashed build. And a test profiles the same 200 rows as parquet and as jsonl -- one folded, one materialised -- and asserts the stats, features and classification match, so the batch size cannot leak into the numbers. Three tests had been passing for the wrong reason: they monkeypatched `measure_columns` to force a measurement failure, which the fold path does not call. One of them still passed because its poisoned partition classified as `unknown` anyway. All three now patch `classify`, which both paths go through. Also caught by ruff and ty rather than by tests: `get_reader` had ended up outside the per-file guard, so a format with no registered reader would have aborted the partition instead of being reported as a FileError. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/pipeline.py | 177 ++++++++++++++--- .../profiler/readers/base.py | 25 +++ .../profiler/readers/jsonl.py | 71 +++++-- .../profiler/readers/parquet.py | 32 +++- .../nemo_datasets_plugin/profiler/stats.py | 179 +++++++++++++----- plugins/nemo-datasets/tests/test_pipeline.py | 79 +++++++- plugins/nemo-datasets/tests/test_stats.py | 4 +- 7 files changed, 468 insertions(+), 99 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index ff2428d25c..402e845a89 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -10,10 +10,17 @@ ``classification`` (roles, format, prompt form, dataset type, and verifiability). Every file is opened — sampling a *subset of files* would hide columns that appear only in later -shards — but a partition's ``row_budget`` is divided across its files, so peak memory tracks the -budget rather than the shard count. Capping each file instead put the knob on the wrong axis: -resharding the same data then multiplied the rows held in memory without describing any more of it. -See :data:`DEFAULT_ROW_BUDGET`. +shards. + +A partition whose files all declare a schema is **folded**: the footers are read first, which gives +the columns and the exact row count before a single row is parsed, and batches are then measured and +let go. Nothing grows with the file, so an exhaustive read costs what a budgeted one costs — +measured at 10.4 MB against 65.1 MB for the same 21,362 rows materialised. ``row_budget`` is a limit +on *work*, not the memory guard it used to be. + +A partition whose files do not all declare a schema is **materialised**, as before: the schema has +to be inferred from the rows, so the rows have to be kept until it has been. That is line-delimited +formats, and folding them needs accumulators created lazily as columns appear, which is still to do. The budget is a target rather than a ceiling. :data:`MIN_ROWS_PER_FILE` is the floor every file is read to however thin its share gets, since one sampled below it cannot contribute the columns it @@ -32,10 +39,15 @@ from nemo_datasets_plugin.profiler.classify import PrefixPairFold, classify from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource from nemo_datasets_plugin.profiler.partition import group_partitions -from nemo_datasets_plugin.profiler.readers.base import detect_format, get_reader, is_unsupported_data +from nemo_datasets_plugin.profiler.readers.base import ( + FilePreview, + detect_format, + get_reader, + is_unsupported_data, +) from nemo_datasets_plugin.profiler.schema import derive_features from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits -from nemo_datasets_plugin.profiler.stats import measure_columns, quote_enumerations +from nemo_datasets_plugin.profiler.stats import ColumnFold, measure_columns, quote_enumerations from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, DatasetProfile, @@ -279,6 +291,81 @@ class _PartitionOutcome: file_errors: list[FileError] # files this partition grouped but could not fully read +def _peek_files(source: FileSource, entries: list[FileEntry]) -> dict[str, FilePreview]: + """What each file declares about itself, before any of them is read. + + A failure here is not reported: it will surface as a :class:`FileError` when the file is actually + read, with a reason, and reporting it twice would double-count. All this decides is whether the + partition can be folded, and a file that cannot be peeked cannot. + """ + previews: dict[str, FilePreview] = {} + for entry in entries: + try: + previews[entry.path] = get_reader(_format_of(entry.path)).peek(source, entry) + except Exception: + previews[entry.path] = FilePreview() + return previews + + +def _expected_rows(previews: dict[str, FilePreview], row_cap: int | None) -> int | None: + """How many rows the fold is about to see, if every file said. + + Capped per file the same way the read will be, so a budgeted run strides its quality sample over + what it will actually scan rather than over what the dataset holds. + """ + total = 0 + for preview in previews.values(): + if preview.num_rows is None: + return None + total += min(preview.num_rows, row_cap) if row_cap is not None else preview.num_rows + return total + + +class _PartitionFolds: + """The two folds a partition needs, driven together over the same batches. + + One is per column; the other compares two columns of a row against each other and so belongs to + neither. Keeping them side by side is what lets the file loop hand over a batch and forget it. + """ + + def __init__(self, features: list[FeatureSchema], expected_rows: int | None) -> None: + self.features = features + self._columns = ColumnFold(features, expected_rows) + self._prefix = PrefixPairFold(features) + + def update(self, rows: list[dict]) -> None: + self._columns.update(rows) + self._prefix.update(rows) + + def measure( + self, column_roles: dict[str, str] + ) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: + """Schema, stats and classification from what was folded. Guarded like :func:`_measure`. + + The columns have their own per-column guard inside the fold; this is the wide one, for + anything structural that no single column owns. + """ + try: + measured = self._columns.finalize() + classification = classify( + self.features, + measured.stats, + probes=measured.probes, + prefix_pair=self._prefix.result(), + column_roles=column_roles, + ) + quote_enumerations(self.features, measured.stats, measured.vocabularies) + classification.evidence.extend(measured.errors) + return self.features, measured.stats, classification + except Exception as exc: + detail = f"could not measure this partition: {type(exc).__name__}: {exc}" + return ( + [], + {}, + PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]), + ) + + def _profile_partition( source: FileSource, name: str, @@ -300,9 +387,27 @@ def _profile_partition( :class:`FileError` the envelope collects, contributes no rows, and flips ``scanned_all`` off — it never aborts the profile. Files that read cleanly are counted, not listed. """ + # Footers first, before a single row is read. A parquet file declares its schema and its exact + # row count there, so one seek per file establishes the partition's whole shape: what the columns + # are, and how many rows are coming. That is what a fold needs and cannot otherwise have -- the + # accumulators must exist before the first batch, and the quality stride must be placed before + # the column it strides has been seen. + previews = _peek_files(source, entries) + can_fold = all(preview is not None and preview.arrow_schema is not None for preview in previews.values()) + partition_rows: list[dict] = [] - arrow_schemas: list[pa.Schema] = [] + arrow_schemas: list[pa.Schema] = [ + preview.arrow_schema for preview in previews.values() if preview is not None and preview.arrow_schema + ] all_declared = True # every file that contributed rows carried a declared schema + folds: _PartitionFolds | None = None + if can_fold: + declared = _unify_schemas(arrow_schemas) + if declared is not None: + folds = _PartitionFolds( + derive_features([], declared), + expected_rows=_expected_rows(previews, _per_file_cap(row_budget, len(entries))), + ) rows_scanned = 0 files_read = 0 rows_present: int | None = 0 @@ -318,32 +423,45 @@ def _profile_partition( for entry in split.entries: file_formats.add(_format_of(entry.path)) error: str | None = None + num_rows: int | None = None + scanned_all = False try: - result = get_reader(_format_of(entry.path)).read(source, entry, row_cap=row_cap) + # Inside the guard: resolving the reader can fail too, and a format with no reader + # registered is a file the profiler could not use like any other. + reader = get_reader(_format_of(entry.path)) + if folds is not None: + # Streamed: batches are folded and let go, so nothing here grows with the file. + num_rows = previews[entry.path].num_rows # from the footer, read before any row + scanned = 0 + for batch in reader.batches(source, entry, row_cap=row_cap): + folds.update(batch) + scanned += len(batch) + files_read += 1 + rows_scanned += scanned + scanned_all = num_rows is not None and scanned >= num_rows + else: + # Materialised: no declared schema, so the schema has to be inferred from the rows + # and the rows have to be kept until it has been. Phase 4 folds this path too. + result = reader.read(source, entry, row_cap=row_cap) + files_read += 1 + error = result.error + num_rows = result.num_rows + rows_scanned += result.rows_scanned + partition_rows.extend(result.rows) + if result.arrow_schema is None and result.rows: + # Rows with no schema behind them: the unified schema no longer covers the + # partition, so _measure must infer from rows rather than trust a partial one. + all_declared = False + # Exhaustive requires parsing every row; a known footer count alone is not enough, + # and a partial read (corrupt lines skipped) is not exhaustive however many it got. + scanned_all = num_rows is not None and result.rows_scanned >= num_rows and error is None except Exception as exc: # Failure isolation: an unreadable file (or missing reader) keeps its identity, # skips its rows, and does not abort the profile. The reason is recorded rather than # swallowed, so a consumer can tell corrupt input from a profiler bug. - result = None error = f"{type(exc).__name__}: {exc}" - if result is None: num_rows = None scanned_all = False - else: - files_read += 1 - error = result.error - num_rows = result.num_rows - rows_scanned += result.rows_scanned - partition_rows.extend(result.rows) - if result.arrow_schema is not None: - arrow_schemas.append(result.arrow_schema) - elif result.rows: - # Rows with no schema behind them: the unified schema no longer covers the - # partition, so _measure must infer from rows rather than trust a partial one. - all_declared = False - # Exhaustive requires parsing every row; a known footer count alone is not enough, and - # a partial read (corrupt lines skipped) is not exhaustive however many rows it got. - scanned_all = num_rows is not None and result.rows_scanned >= num_rows and error is None if error is not None: file_errors.append(FileError(path=entry.path, error=error)) rows_present = _add_known(rows_present, num_rows) @@ -368,9 +486,12 @@ def _profile_partition( num_examples=split_examples if split_counts_known else None, ) ) - features, stats, classification = _measure( - partition_rows, arrow_schemas, all_declared=all_declared, column_roles=column_roles - ) + if folds is not None: + features, stats, classification = folds.measure(column_roles) + else: + features, stats, classification = _measure( + partition_rows, arrow_schemas, all_declared=all_declared, column_roles=column_roles + ) partition = PartitionProfile( name=name, # Observed, not chosen: the partition reports the formats its files turned out to be in diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py index 4afb575a80..796bec7728 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py @@ -12,6 +12,7 @@ from __future__ import annotations +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from typing import Any, ClassVar, Protocol @@ -34,15 +35,39 @@ class ReadResult: error: str | None = None +@dataclass(frozen=True) +class FilePreview: + """What a reader can learn about a file without reading a single row. + + A parquet footer carries both; a line-delimited format carries neither. The pipeline asks every + file this before it reads any of them, because knowing the schema up front is what lets a + partition be measured without first being materialised, and knowing the row count up front is + what lets a quality sample be strided across a column the fold has not finished seeing. + """ + + arrow_schema: pa.Schema | None = None + num_rows: int | None = None + + class FormatReader(Protocol): """Reads schema and rows for one file format.""" file_format: ClassVar[str] + def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: + """What the file declares about itself, without reading its rows.""" + ... + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: """Read up to ``row_cap`` rows (all rows when None) plus whatever the format declares cheaply.""" ... + def batches( + self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None + ) -> Iterator[list[dict[str, Any]]]: + """The same rows :meth:`read` would return, handed over in chunks and never all at once.""" + ... + _READERS: dict[str, FormatReader] = {} _builtins_loaded = False diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py index 85a12395a0..8ce2b61672 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py @@ -6,35 +6,80 @@ from __future__ import annotations import json +from collections.abc import Iterator from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource -from nemo_datasets_plugin.profiler.readers.base import ReadResult, register_reader +from nemo_datasets_plugin.profiler.readers.base import FilePreview, ReadResult, register_reader + +# Rows handed over at a time by :meth:`JsonlReader.batches`, matching the parquet reader so the +# caller's working set does not depend on which format it happens to be folding. +_BATCH_ROWS = 1024 + + +def _records(stream) -> Iterator[tuple[dict | None, str | None]]: + """Each line of the stream as either a record or a reason it was not one. + + Shared by :meth:`JsonlReader.read` and :meth:`JsonlReader.batches` so the two cannot drift on + what counts as a row -- a blank line, a stray scalar and a truncated line are three different + things and only one of them is an error. + """ + for line_number, raw_line in enumerate(stream, start=1): + stripped = raw_line.strip() + if not stripped: # tolerate blank lines between records + continue + try: + record = json.loads(stripped) + except ValueError as exc: + # A truncated or corrupt line costs that line, never the file. Dropping the whole file + # would erase its row count and every column it was the only witness for. + yield None, f"line {line_number}: {exc}" + continue + if isinstance(record, dict): + yield record, None + # a record is a column map; stray scalars/arrays are skipped rather than crash downstream class JsonlReader: file_format = "jsonl" + def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: + """Nothing. A line-delimited file declares no schema and carries no row count, which is why + a partition holding one cannot be folded without reading it first.""" + return FilePreview() + + def batches(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> Iterator[list[dict]]: + """Rows in chunks. Parse errors are silent here -- :meth:`read` is the path that accounts for + them, and the fold path does not reach a file with no declared schema.""" + rows: list[dict] = [] + scanned = 0 + with source.open(entry.path) as stream: + for record, _ in _records(stream): + if record is None: + continue + rows.append(record) + scanned += 1 + if len(rows) >= _BATCH_ROWS: + yield rows + rows = [] + if row_cap is not None and scanned >= row_cap: + break + if rows: + yield rows + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: rows: list[dict] = [] unparseable = 0 first_failure: str | None = None hit_cap = False with source.open(entry.path) as stream: - for line_number, raw_line in enumerate(stream, start=1): - stripped = raw_line.strip() - if not stripped: # tolerate blank lines between records - continue - try: - record = json.loads(stripped) - except ValueError as exc: - # A truncated or corrupt line costs that line, never the file. Dropping the whole - # file would erase its row count and every column it was the only witness for. + for record, failure in _records(stream): + # Branching on the record rather than the failure: the two are exclusive, and this + # way the type narrows without an ignore standing in for the reasoning. + if record is None: unparseable += 1 if first_failure is None: - first_failure = f"line {line_number}: {exc}" + first_failure = failure continue - if not isinstance(record, dict): - continue # a record is a column map; skip stray scalars/arrays rather than crash downstream rows.append(record) if row_cap is not None and len(rows) >= row_cap: hit_cap = True diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py index bea95ce212..62863bf682 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py @@ -5,14 +5,44 @@ from __future__ import annotations +from collections.abc import Iterator + import pyarrow.parquet as pq from nemo_datasets_plugin.profiler.file_source import FileEntry, FileSource -from nemo_datasets_plugin.profiler.readers.base import ReadResult, register_reader +from nemo_datasets_plugin.profiler.readers.base import FilePreview, ReadResult, register_reader + +# Rows handed over at a time. Small enough that the working set is a knob independent of the +# dataset, large enough that per-batch overhead stays invisible. +_BATCH_ROWS = 1024 class ParquetReader: file_format = "parquet" + def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: + """Schema and exact row count from the footer. Reads no rows, so a partition's whole shape is + knowable for the cost of one seek per file.""" + with source.open(entry.path) as stream: + parquet_file = pq.ParquetFile(stream) + return FilePreview(arrow_schema=parquet_file.schema_arrow, num_rows=parquet_file.metadata.num_rows) + + def batches(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> Iterator[list[dict]]: + """Rows in chunks, so the caller can fold them and let each chunk go.""" + scanned = 0 + with source.open(entry.path) as stream: + parquet_file = pq.ParquetFile(stream) + if row_cap == 0: + return + for batch in parquet_file.iter_batches(batch_size=min(row_cap or _BATCH_ROWS, _BATCH_ROWS)): + rows = batch.to_pylist() + if row_cap is not None and scanned + len(rows) > row_cap: + rows = rows[: row_cap - scanned] + scanned += len(rows) + if rows: + yield rows + if row_cap is not None and scanned >= row_cap: + return + def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: with source.open(entry.path) as stream: parquet_file = pq.ParquetFile(stream) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index f9d0010a67..07ede3d71c 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -85,58 +85,96 @@ class ColumnMeasurements: errors: list[Evidence] -def measure_columns(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> ColumnMeasurements: - """Measure every top-level column: its statistics and its content probes, in one pass each. +class ColumnFold: + """The per-column accumulators for one partition, fed batch by batch. Each column is isolated. A value no detector anticipated -- a chat message whose ``role`` is a number, a float where a string was declared -- costs that column its measurements and nothing else, where previously it cost the partition every measurement it had. The failure is reported as an ``error`` evidence rather than left as a silent gap, because a column absent from ``stats`` - is otherwise indistinguishable from one that simply had nothing worth measuring. + is otherwise indistinguishable from one that simply had nothing worth measuring. It is caught per + column *per batch*, so a bad row in the middle of a file cannot take the rest of the file with it. This is the narrow half of the two guards the profiler runs. The wide one still wraps the whole measure stage, and still catches anything structural -- schema derivation, classification -- that is not attributable to a single column. - Statistics and probes are measured together because they read the same values, and extracting a - column out of the rows costs more than either measurement. Neither fills in + Statistics and probes are folded together because they read the same values, and extracting a + column out of a batch costs more than either measurement. Neither fills in ``categorical.values``: that needs the roles, which classification has not assigned yet, so - :func:`quote_enumerations` adds them afterwards. + :func:`quote_enumerations` adds them afterwards, from the vocabulary this kept. """ - stats: dict[str, ColumnStats] = {} - probes: dict[str, ColumnProbes] = {} - vocabularies: dict[str, set[Any]] = {} - errors: list[Evidence] = [] - for feature in features: - # Parquet permits duplicate field names, and both maps are keyed by name. Measuring the - # first and skipping the rest makes which one wins deterministic instead of "whichever came - # last", and keeps stats and probes agreeing on the same one. - if feature.name in probes: - continue - accumulator = _accumulator_for(feature) - try: - # One batch, being every row this partition holds. The accumulator does not care: fed in - # pieces it gives the same answer, which is what lets the caller stop materialising the - # partition first. - accumulator.update([row.get(feature.name) for row in rows]) - column, probe = accumulator.finalize() - probes[feature.name] = probe - vocabulary = accumulator.vocabulary() - if vocabulary is not None: - vocabularies[feature.name] = vocabulary - except Exception as exc: - errors.append( - Evidence( + + def __init__(self, features: list[FeatureSchema], expected_rows: int | None = None) -> None: + self._accumulators: dict[str, ColumnAccumulator] = {} + self._features: list[FeatureSchema] = [] + self._failed: dict[str, Evidence] = {} + for feature in features: + # Parquet permits duplicate field names, and every map here is keyed by name. Measuring + # the first and skipping the rest makes which one wins deterministic instead of + # "whichever came last", and keeps stats and probes agreeing on the same one. + if feature.name in self._accumulators: + continue + self._accumulators[feature.name] = _accumulator_for(feature, expected_rows) + self._features.append(feature) + + def update(self, rows: list[dict[str, Any]]) -> None: + """Fold one batch of rows into every column's accumulator.""" + for feature in self._features: + if feature.name in self._failed: + continue + try: + self._accumulators[feature.name].update([row.get(feature.name) for row in rows]) + except Exception as exc: + self._failed[feature.name] = Evidence( kind="error", detail=( f"column {feature.name!r} ({feature.dtype}) could not be measured: {type(exc).__name__}: {exc}" ), ) - ) - continue - if column is not None: - stats[feature.name] = column - return ColumnMeasurements(stats=stats, probes=probes, vocabularies=vocabularies, errors=errors) + + def finalize(self) -> ColumnMeasurements: + stats: dict[str, ColumnStats] = {} + probes: dict[str, ColumnProbes] = {} + vocabularies: dict[str, set[Any]] = {} + errors: list[Evidence] = [] + for feature in self._features: + failure = self._failed.get(feature.name) + if failure is not None: + errors.append(failure) + continue + accumulator = self._accumulators[feature.name] + try: + column, probe = accumulator.finalize() + except Exception as exc: + errors.append( + Evidence( + kind="error", + detail=( + f"column {feature.name!r} ({feature.dtype}) could not be summarised: " + f"{type(exc).__name__}: {exc}" + ), + ) + ) + continue + probes[feature.name] = probe + vocabulary = accumulator.vocabulary() + if vocabulary is not None: + vocabularies[feature.name] = vocabulary + if column is not None: + stats[feature.name] = column + return ColumnMeasurements(stats=stats, probes=probes, vocabularies=vocabularies, errors=errors) + + +def measure_columns(features: list[FeatureSchema], rows: list[dict[str, Any]]) -> ColumnMeasurements: + """Measure every top-level column over rows already in hand. + + The whole partition as a single batch. :class:`ColumnFold` is the same measurement taken as the + rows arrive; this is the shape for a caller that has them all anyway. + """ + fold = ColumnFold(features) + fold.update(rows) + return fold.finalize() def quote_enumerations( @@ -301,29 +339,36 @@ def values(self) -> set[Any] | None: class StringAccumulator(ColumnAccumulator): """A ``string`` column: length quantiles, corruption ratios, and a vocabulary if it has one.""" - def __init__(self) -> None: + def __init__(self, expected_rows: int | None = None) -> None: super().__init__() - self._strings: list[str] = [] self._lengths = _LengthHistogram() self._vocabulary = _Vocabulary() + # With the row count known -- parquet footers give it before a row is read -- the quality + # stride can be placed now and each string measured or skipped as it goes by, retaining + # nothing. Without it there is no stride to compute yet, so the strings are held and strided + # at the end. That is the one term here still sized by the column, and it is exactly the + # partitions that have no footer to read. + self._stride = _quality_stride(expected_rows) if expected_rows is not None else None + self._quality = _TextQualityCounters() + self._strings: list[str] = [] + self._seen = 0 def _observe(self, present: list[Any]) -> None: for value in present: if isinstance(value, str): - # The length folds away; the string itself is retained only because the quality - # stride needs the column's row count to place its sample, and that is not known - # until the last batch. Bounding that is the parquet footer sum's job -- it is the - # one term here still sized by the column, and it costs what materialising the - # column already cost. - self._strings.append(value) self._lengths.add(len(value)) + if self._stride is None: + self._strings.append(value) + elif self._seen % self._stride == 0: + self._quality.add(value) + self._seen += 1 self._vocabulary.update(present) def _blocks(self) -> dict[str, Any]: text = quality = None - if self._strings: + if self._seen: text = TextStats(chars=self._lengths.quantiles()) - quality = _text_quality(self._strings) + quality = _text_quality(self._strings) if self._stride is None else self._quality.finalize() return {"text": text, "quality": quality, "categorical": self._vocabulary.finalize()} def vocabulary(self) -> set[Any] | None: @@ -440,10 +485,14 @@ def _blocks(self) -> dict[str, Any]: } -def _accumulator_for(feature: FeatureSchema) -> ColumnAccumulator: - """The accumulator that knows how to measure this column, dispatched once on its dtype.""" +def _accumulator_for(feature: FeatureSchema, expected_rows: int | None = None) -> ColumnAccumulator: + """The accumulator that knows how to measure this column, dispatched once on its dtype. + + ``expected_rows`` is the partition's row count when it is known before reading -- only a string + column uses it, to place its quality stride without retaining the column. + """ if feature.dtype == "string": - return StringAccumulator() + return StringAccumulator(expected_rows) if feature.dtype == "bool": return BoolAccumulator() if feature.dtype == "messages": @@ -607,6 +656,40 @@ def _non_ascii_count(text: str) -> int: return _count_matches(_NON_ASCII_RUN, text) +class _TextQualityCounters: + """The three corruption ratios as running sums, so a strided sample needs no storage. + + Every denominator is the sample's own, never the column's: each ratio is an estimate over the + rows actually scanned, which is what keeps it unbiased rather than diluted. + """ + + def __init__(self) -> None: + self._chars = 0 + self._whitespace = 0 + self._non_ascii = 0 + self._repetition = 0.0 + self._rows = 0 + + def add(self, text: str) -> None: + self._chars += len(text) + self._whitespace += _whitespace_count(text) + self._non_ascii += _non_ascii_count(text) + self._repetition += _repetition_score(text) + self._rows += 1 + + def finalize(self) -> TextQuality: + return TextQuality( + whitespace_ratio=self._whitespace / self._chars if self._chars else 0.0, + non_ascii_ratio=self._non_ascii / self._chars if self._chars else 0.0, + repetition_score=self._repetition / self._rows if self._rows else 0.0, + ) + + +def _quality_stride(rows: int) -> int: + """How many rows to step between quality samples, given the column's length.""" + return max(1, rows // _QUALITY_SAMPLE_ROWS) + + def _text_quality(strings: list[str]) -> TextQuality: sample = _quality_sample(strings) total_chars = 0 diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 648ba4efb9..b2911cfa8d 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -12,7 +12,8 @@ import pytest from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource from nemo_datasets_plugin.profiler.partition import group_partitions -from nemo_datasets_plugin.profiler.pipeline import _measure, profile +from nemo_datasets_plugin.profiler.pipeline import _expected_rows, _measure, _peek_files, profile +from nemo_datasets_plugin.profiler.readers.base import FilePreview from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits from nemo_platform_plugin.files.dataset_profile import DatasetProfile @@ -158,6 +159,66 @@ def test_data_files_glob_means_the_same_thing_to_pythons_own_glob(tmp_path): assert all(name.startswith(f"data/{split.name}") for name in resolved) +# --- the fold ------------------------------------------------------------------------------------ + + +def test_a_parquet_footer_declares_enough_to_fold_without_reading_rows(tmp_path): + # The footer is what makes a fold possible at all: the schema, so accumulators can exist before + # the first batch, and the exact row count, so the quality stride can be placed before the column + # it strides has been seen. A line-delimited file declares neither, which is why it materialises. + _write_parquet(tmp_path / "train.parquet", [{"a": i} for i in range(7)]) + (tmp_path / "extra.jsonl").write_text('{"a": 1}\n') + source = LocalFileSource(tmp_path) + + previews = _peek_files(source, source.list_files()) + + assert previews["train.parquet"].num_rows == 7 + assert previews["train.parquet"].arrow_schema is not None + assert previews["extra.jsonl"] == FilePreview() # declares nothing, so the partition cannot fold + + +def test_expected_rows_counts_what_the_read_will_actually_scan(): + # The stride has to be placed over the rows that will be *scanned*, not the rows the dataset + # holds, or a budgeted run would stride far too coarsely and sample almost nothing. + previews = {"a": FilePreview(num_rows=100), "b": FilePreview(num_rows=100)} + assert _expected_rows(previews, None) == 200 + assert _expected_rows(previews, 30) == 60 # capped per file, exactly as the read will be + assert _expected_rows({"a": FilePreview(num_rows=100), "b": FilePreview()}, None) is None + + +def test_the_folded_and_materialised_paths_measure_the_same_thing(tmp_path): + # Parquet declares a schema and is folded batch by batch; jsonl declares none and is + # materialised. The same rows have to measure the same either way, or the batch size -- an + # implementation detail no reader of a profile can see -- would be visible in the numbers. + rows = [{"prompt": f"question {i}", "completion": "answer " * (i % 7 + 1), "score": i % 5} for i in range(200)] + _write_parquet(tmp_path / "pq" / "train.parquet", rows) + (tmp_path / "jl").mkdir() + (tmp_path / "jl" / "train.jsonl").write_text("\n".join(json.dumps(row) for row in rows)) + + folded = profile(LocalFileSource(tmp_path / "pq"), created_at=FIXED_TIME).partitions[0] + materialised = profile(LocalFileSource(tmp_path / "jl"), created_at=FIXED_TIME).partitions[0] + + assert folded.stats == materialised.stats + assert folded.classification == materialised.classification + assert [f.model_dump() for f in folded.features] == [f.model_dump() for f in materialised.features] + + +def test_an_exhaustive_fold_does_not_cost_more_than_a_budgeted_one(tmp_path): + # The point of the whole exercise: reading every row costs what reading some of them costs, so + # the budget stops being a memory guard. Same measurements, and `stats_complete` finally true. + _write_parquet(tmp_path / "train.parquet", [{"t": f"row {i}" * (i % 5 + 1)} for i in range(5000)]) + + budgeted = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=500) + exhaustive = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=None) + + assert budgeted.sampling.rows_scanned == 500 + assert exhaustive.sampling.rows_scanned == 5000 + assert budgeted.partitions[0].stats_complete is False + assert exhaustive.partitions[0].stats_complete is True + # Exact where it claims to be exact: the longest row is found by reading all of them. + assert exhaustive.partitions[0].stats["t"].text.chars.max >= budgeted.partitions[0].stats["t"].text.chars.max + + # --- partition grouping -------------------------------------------------------------------------- @@ -529,7 +590,7 @@ def test_profile_degrades_one_partition_when_measurement_fails(tmp_path, monkeyp from nemo_datasets_plugin.profiler import pipeline as pipeline_module _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) - monkeypatch.setattr(pipeline_module, "measure_columns", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + monkeypatch.setattr(pipeline_module, "classify", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise @@ -561,7 +622,7 @@ def test_a_measurement_failure_does_not_look_like_a_read_failure(tmp_path, monke from nemo_datasets_plugin.profiler import pipeline as pipeline_module _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}]) - monkeypatch.setattr(pipeline_module, "measure_columns", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + monkeypatch.setattr(pipeline_module, "classify", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) @@ -587,7 +648,9 @@ def _observe(self, present): monkeypatch.setattr( stats_module, "_accumulator_for", - lambda feature: Boom() if feature.name == "completion" else real_accumulator_for(feature), + lambda feature, expected_rows=None: ( + Boom() if feature.name == "completion" else real_accumulator_for(feature, expected_rows) + ), ) _write_parquet(tmp_path / "train.parquet", [{"prompt": "q", "completion": "a"}]) @@ -603,16 +666,16 @@ def _observe(self, present): def test_a_measurement_failure_is_scoped_to_its_own_partition(tmp_path, monkeypatch): from nemo_datasets_plugin.profiler import pipeline as pipeline_module - real_measure_columns = pipeline_module.measure_columns + real_classify = pipeline_module.classify - def poison_one_partition(features, rows): + def poison_one_partition(features, stats, **kwargs): if any(feature.name == "poison" for feature in features): raise RuntimeError("boom") - return real_measure_columns(features, rows) + return real_classify(features, stats, **kwargs) _write_parquet(tmp_path / "good" / "train.parquet", [{"prompt": "q", "completion": "a"}]) _write_parquet(tmp_path / "bad" / "train.parquet", [{"poison": 1}]) - monkeypatch.setattr(pipeline_module, "measure_columns", poison_one_partition) + monkeypatch.setattr(pipeline_module, "classify", poison_one_partition) partitions = {p.name: p for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions} diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index e6f4e66b53..4fc83f7b84 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -117,7 +117,9 @@ def _observe(self, present): monkeypatch.setattr( stats_module, "_accumulator_for", - lambda feature: Boom() if feature.name == "bad" else real_accumulator_for(feature), + lambda feature, expected_rows=None: ( + Boom() if feature.name == "bad" else real_accumulator_for(feature, expected_rows) + ), ) features = [_feature("good", "string"), _feature("bad", "string")] From 65acf684fe8a4209513b6c3621fb9d854af79699 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 19:58:04 -0400 Subject: [PATCH 41/44] fix(datasets): bound the two things row content could grow without limit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4, first step: §5 of the streaming spec. Two structures were fed straight from row content with nothing stopping them, and neither needed a bound before because the row budget was one by accident. `_features_from_rows` builds a key union across every row, so a malformed file whose rows carry unique keys mints a column per row -- and since Phase 3, an accumulator per column with it. `MAX_COLUMNS = 4096` stops that. The truncation is reported rather than silent: a profile that described 4,096 of a file's columns as though they were all of them is worse than one that failed, because the reader has no way to tell a wide table from a broken one. Applied to declared schemas too, which are bounded by their file but not by anything sensible. `roles_seen` is worse in one respect: membership is checked against the list, so an unbounded one is quadratic as well as unbounded. `_MAX_ROLES_SEEN = 64`. That truncation is silent, and deliberately so -- the list exists so a reader can pick a chat template, and a column with more than sixty-four distinct roles is not a chat column, which the first few dozen already say. The contract says it is bounded. Reachable today, not hypothetically: with `row_budget=None` there is nothing else holding either of them down, and Phase 3 made an unbounded read the affordable default rather than the expensive exception. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 4 ++- .../nemo_datasets_plugin/profiler/pipeline.py | 23 +++++++++++++++- .../nemo_datasets_plugin/profiler/schema.py | 26 +++++++++++++++++-- .../nemo_datasets_plugin/profiler/stats.py | 9 ++++++- plugins/nemo-datasets/tests/test_schema.py | 14 +++++++++- plugins/nemo-datasets/tests/test_stats.py | 10 +++++++ 6 files changed, 80 insertions(+), 6 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 2b83760a4e..295445c854 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -174,7 +174,9 @@ class MessageStats(BaseModel): '"user", "assistant", "tool"], but equally ShareGPT\'s ["human", "gpt"] or a house convention. ' "A measurement of row content, not a vocabulary the profiler picks from, so it is deliberately " "not an enum: an unexpected role is the finding worth reporting, and normalizing or dropping it " - "would hide exactly what a consumer needs to see before choosing a chat template." + "would hide exactly what a consumer needs to see before choosing a chat template. " + "Bounded: this is fed straight from row content, and a column with more distinct roles " + "than fit here is not a chat column, which the first few dozen already say." ), ) ends_with_assistant_rate: float = Field( diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 402e845a89..bcfd98bfc7 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -45,7 +45,7 @@ get_reader, is_unsupported_data, ) -from nemo_datasets_plugin.profiler.schema import derive_features +from nemo_datasets_plugin.profiler.schema import MAX_COLUMNS, columns_were_capped, derive_features from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits from nemo_datasets_plugin.profiler.stats import ColumnFold, measure_columns, quote_enumerations from nemo_platform_plugin.files.dataset_profile import ( @@ -273,6 +273,7 @@ def _measure( quote_enumerations(features, measured.stats, measured.vocabularies) # After classify, so its own reasoning reads first and a column that could not be measured # is reported as a caveat on the result rather than as part of the case for it. + classification.evidence.extend(_capped_columns_evidence(features)) classification.evidence.extend(measured.errors) return features, measured.stats, classification except Exception as exc: @@ -291,6 +292,25 @@ class _PartitionOutcome: file_errors: list[FileError] # files this partition grouped but could not fully read +def _capped_columns_evidence(features: list[FeatureSchema]) -> list[Evidence]: + """Say so when the schema stopped at the cap rather than at the end of the data. + + A profile that quietly described 4,096 of a file's columns as though they were all of them would + be worse than one that failed: the reader has no way to tell a wide table from a broken one. + """ + if not columns_were_capped(features): + return [] + return [ + Evidence( + kind="error", + detail=( + f"stopped at {MAX_COLUMNS} columns; the rest of this partition's schema was not " + f"described. A file whose rows carry unique keys will do this." + ), + ) + ] + + def _peek_files(source: FileSource, entries: list[FileEntry]) -> dict[str, FilePreview]: """What each file declares about itself, before any of them is read. @@ -355,6 +375,7 @@ def measure( column_roles=column_roles, ) quote_enumerations(self.features, measured.stats, measured.vocabularies) + classification.evidence.extend(_capped_columns_evidence(self.features)) classification.evidence.extend(measured.errors) return self.features, measured.stats, classification except Exception as exc: diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py index f5e5ccbb14..51c199cd67 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py @@ -24,16 +24,34 @@ _MESSAGE_KEY_SETS = ({"role", "content"}, {"from", "value"}) +# Columns a partition may have before the profiler stops describing it. Nothing legitimate reaches +# this: it is a guard against a malformed file whose rows carry per-row unique keys, which would +# otherwise mint a column -- and an accumulator -- for every row in the dataset. The row budget used +# to bound this by accident, since a schema inferred from at most N rows had at most N keys; an +# unbounded read has no such accident, so the bound is stated. +MAX_COLUMNS = 4096 + + def derive_features(rows: list[dict[str, Any]], arrow_schema: pa.Schema | None = None) -> list[FeatureSchema]: - """The row schema. Uses the declared arrow schema when present, else infers from ``rows``.""" + """The row schema. Uses the declared arrow schema when present, else infers from ``rows``. + + Truncated at :data:`MAX_COLUMNS`. Use :func:`columns_were_capped` to tell a dataset that really + is that wide from one whose keys are runaway; the caller reports the difference rather than + silently describing part of a file as though it were all of it. + """ if arrow_schema is not None: return [ _feature_from_arrow(arrow_schema.field(i).name, arrow_schema.field(i).type) - for i in range(len(arrow_schema)) + for i in range(min(len(arrow_schema), MAX_COLUMNS)) ] return _features_from_rows(rows) +def columns_were_capped(features: list[FeatureSchema]) -> bool: + """Whether the schema stopped at the cap rather than at the end of the data.""" + return len(features) >= MAX_COLUMNS + + def _is_message_struct(item: FeatureSchema) -> bool: if item.dtype != "struct" or item.fields is None: return False @@ -93,6 +111,10 @@ def _features_from_rows(rows: list[dict[str, Any]]) -> list[FeatureSchema]: if key not in seen: seen.add(key) ordered_keys.append(key) + if len(ordered_keys) >= MAX_COLUMNS: + break + if len(ordered_keys) >= MAX_COLUMNS: + break return [_infer_feature(key, [row.get(key) for row in rows]) for key in ordered_keys] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 07ede3d71c..6a108f67f5 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -457,7 +457,7 @@ def _observe(self, present: list[Any]) -> None: # that an unexpected role is the finding worth surfacing, not something to # normalize away. role = role if isinstance(role, str) else str(role) - if role not in self._roles_seen: + if role not in self._roles_seen and len(self._roles_seen) < _MAX_ROLES_SEEN: self._roles_seen.append(role) total_content += _content_len(_message_field(message, "content", "value")) # `.get` truthiness, not `in`: parquet materializes every declared struct field, so a @@ -734,6 +734,13 @@ def _repetition_score(text: str) -> float: # target — a false negative over a large slice of public chat data. _ASSISTANT_ROLES = {"assistant", "gpt", "bot", "model", "chatbot", "ai"} +# Distinct role strings a chat column may show before the list stops growing. It is fed straight from +# row content, so without a bound one malformed column could hold a string per message -- and since +# membership is checked against the list, that is quadratic as well as unbounded. The truncation +# costs nothing a reader would act on: the list exists to pick a chat template, and a column with +# more than this many roles is not a chat column, which the first few dozen already say. +_MAX_ROLES_SEEN = 64 + def _message_field(message: dict, *names: str) -> Any: """The first present, non-null value among ``names``. diff --git a/plugins/nemo-datasets/tests/test_schema.py b/plugins/nemo-datasets/tests/test_schema.py index e8fd3364fd..d27b84b91d 100644 --- a/plugins/nemo-datasets/tests/test_schema.py +++ b/plugins/nemo-datasets/tests/test_schema.py @@ -4,7 +4,7 @@ """Tests for row-schema derivation (from a declared arrow schema and from sampled rows).""" import pyarrow as pa -from nemo_datasets_plugin.profiler.schema import derive_features +from nemo_datasets_plugin.profiler.schema import MAX_COLUMNS, columns_were_capped, derive_features # --- from a declared arrow schema (parquet) ------------------------------------------------------ @@ -86,3 +86,15 @@ def test_from_rows_all_null_column_is_json(): def test_derive_features_prefers_declared_arrow_schema(): feature = derive_features([{"x": 1}], pa.schema([("x", pa.int32())]))[0] assert feature.dtype == "int32" # declared width beats the int64 inference from rows + + +def test_column_count_is_bounded_and_says_when_it_stopped(): + # A malformed file whose rows carry unique keys would otherwise mint a column -- and later an + # accumulator -- for every row. The row budget used to bound this by accident; an unbounded read + # does not, so the bound is stated and the truncation is reported rather than silent. + rows = [{f"col{i}": i} for i in range(MAX_COLUMNS + 500)] + features = derive_features(rows) + assert len(features) == MAX_COLUMNS + assert columns_were_capped(features) + + assert not columns_were_capped(derive_features([{"a": 1, "b": 2}])) diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 4fc83f7b84..70d567ae28 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -87,6 +87,16 @@ def test_an_empty_histogram_reports_zeros(): assert (quantiles.p50, quantiles.p95, quantiles.p99, quantiles.max) == (0, 0, 0, 0) +def test_roles_seen_stops_growing(): + # Fed straight from row content, so without a bound one malformed column could hold a string per + # message -- and membership is checked against the list, so it is quadratic as well as unbounded. + rows = _rows("m", [[{"role": f"role-{i}", "content": "x"}] for i in range(stats_module._MAX_ROLES_SEEN * 3)]) + measured = _stats([_feature("m", "messages")], rows)["m"] + assert len(measured.messages.roles_seen) == stats_module._MAX_ROLES_SEEN + # The rates still count every row: only the vocabulary of roles is bounded, not the measurement. + assert measured.messages.turns.max == 1 + + # --- text ---------------------------------------------------------------------------------------- From 85bf215c296ad1954177b8ff2d05a43f7466f92f Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 20:36:40 -0400 Subject: [PATCH 42/44] feat(datasets): fold partitions that declare no schema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4. Every partition now folds; nothing is materialised. The blocker was never lazy column creation, which §7 of the spec assumed. It was dtypes. An accumulator is chosen *by* dtype, and for an inferred schema the dtype is a whole-column decision -- observed types unioned, a disagreement widened to `json` -- so the choice cannot be made while making it still matters. Batch 1 says string, batch 50 says int, and a StringAccumulator has been folding the wrong thing for forty-nine batches. Deferring the choice is the only resolution that neither reads the data twice nor decides from a prefix and hopes. `DeferredAccumulator` measures every shape at once and picks the answer at the end. It costs nothing extra per value -- a string only ever reaches the string state -- and what it costs is four bounded structures per column instead of one. Alongside it `SchemaFold`, which is `_infer_feature` written incrementally. That turned out to be transcription rather than invention: the function was already a set union over observed types, a union over a struct's child keys, and a recursion over a list's flattened elements, all of which are state proportional to the schema and not to the row count. Checked against the original across fourteen shapes -- widening, mixed, nested structs, both chat spellings, empty lists -- at four chunkings each, with no mismatches. Columns are created on first sight and back-filled with the rows they were absent for, which is a pair of additions rather than a pass, and is what makes the result identical to inferring the schema first and measuring second: a row without the key genuinely holds a null for it. The quality stride no longer needs a row count either, so a string column retains nothing on any path. With a footer the stride is fixed and the sample spread evenly; without one it starts at one and doubles as the sample fills, with each sampled row weighted by the stride it stood for -- Horvitz-Thompson, so the estimate stays unbiased rather than weighted toward the head where sampling was densest. That was the last O(rows) term anywhere in the fold: a 60,000-row jsonl went from 9.6 MB to 1.0 MB, and now costs less exhaustively than it did budgeted. Two bugs found on the way, neither by the tests: `_resolve_scalar` ended in `dtypes.pop()`, which mutates its argument. Harmless while every caller passed a throwaway set; `SchemaFold` passes the one it is still using, so the first `finalize()` emptied it and the second resolved the same column to `json`. `batches()` had no way to report a line it could not parse, so a partially read jsonl folded silently and looked complete -- `read()` reported it, and the fold path does not call `read()`. A generator cannot return that: by the time it knows, the caller has consumed everything it yielded. It takes an `errors` list instead. `_measure` is deleted, having lost its last caller and been quietly broken since `PrefixPairFold` stopped taking a schema. `PrefixPairFold` now resolves its two columns off each row, which is what lets it run over a partition whose columns are not known yet. The test that covered `_measure`'s inference is now an end-to-end one against the pipeline, where that behaviour actually lives. Parquet profiles are byte-identical to Phase 3's. Signed-off-by: Albert Cui --- .../nemo_datasets_plugin/profiler/classify.py | 36 +-- .../nemo_datasets_plugin/profiler/pipeline.py | 183 +++++--------- .../profiler/readers/base.py | 15 +- .../profiler/readers/jsonl.py | 26 +- .../profiler/readers/parquet.py | 15 +- .../nemo_datasets_plugin/profiler/schema.py | 100 ++++++-- .../nemo_datasets_plugin/profiler/stats.py | 236 +++++++++++++----- plugins/nemo-datasets/tests/test_classify.py | 2 +- plugins/nemo-datasets/tests/test_pipeline.py | 34 ++- plugins/nemo-datasets/tests/test_stats.py | 39 ++- 10 files changed, 442 insertions(+), 244 deletions(-) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py index 311a5061c4..146b7fcd61 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/classify.py @@ -341,6 +341,12 @@ class PrefixPair: shared: int = 0 +# The column names the alias table maps to the two sides of a preference pair. Looked up directly +# rather than resolved through roles, because this fold runs before classification has assigned any. +_CHOSEN_NAMES = tuple(name for name, role in _ALIAS_ROLES.items() if role == "chosen") +_REJECTED_NAMES = tuple(name for name, role in _ALIAS_ROLES.items() if role == "rejected") + + class PrefixPairFold: """The one probe that reads two columns against each other rather than each on its own. @@ -348,31 +354,25 @@ class PrefixPairFold: between them, and no per-column measurement can see that. Being relational, it also cannot live on a column accumulator, so it folds separately over the same batches. - Resolved by column *name*, not by role, because the fold runs before classification has assigned - any roles -- the same inversion the content probes made when they stopped being role-gated. The - names are the ones the alias table maps to these two roles. + It takes no schema. Resolving by name straight off each row is what lets it run over a partition + whose columns are not known yet, which is every partition with no declared schema. Values that + are not text contribute nothing, so a `chosen` column that turns out to hold chat rather than + strings simply never counts a pair -- the same answer the dtype check used to give up front. """ - def __init__(self, features: list[FeatureSchema]) -> None: - self._left = self._column_named(features, "chosen") - self._right = self._column_named(features, "rejected") + def __init__(self) -> None: self._pairs = 0 self._shared = 0 - @staticmethod - def _column_named(features: list[FeatureSchema], role: str) -> str | None: - wanted = {name for name, aliased in _ALIAS_ROLES.items() if aliased == role} - return next((f.name for f in features if f.name in wanted and f.dtype == "string"), None) - def update(self, rows: list[dict]) -> None: - if self._left is None or self._right is None: - return for row in rows: - left, right = row.get(self._left), row.get(self._right) - if isinstance(left, str) and isinstance(right, str): - self._pairs += 1 - if _common_prefix_len(left, right) >= _EMBEDDED_PROMPT_PREFIX_CHARS: - self._shared += 1 + left = next((row.get(name) for name in _CHOSEN_NAMES if isinstance(row.get(name), str)), None) + right = next((row.get(name) for name in _REJECTED_NAMES if isinstance(row.get(name), str)), None) + if left is None or right is None: + continue + self._pairs += 1 + if _common_prefix_len(left, right) >= _EMBEDDED_PROMPT_PREFIX_CHARS: + self._shared += 1 def result(self) -> PrefixPair: return PrefixPair(pairs=self._pairs, shared=self._shared) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index bcfd98bfc7..4742f7fcc6 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -12,15 +12,18 @@ Every file is opened — sampling a *subset of files* would hide columns that appear only in later shards. -A partition whose files all declare a schema is **folded**: the footers are read first, which gives -the columns and the exact row count before a single row is parsed, and batches are then measured and -let go. Nothing grows with the file, so an exhaustive read costs what a budgeted one costs — -measured at 10.4 MB against 65.1 MB for the same 21,362 rows materialised. ``row_budget`` is a limit -on *work*, not the memory guard it used to be. +Every partition is **folded**: batches are measured and let go, and nothing kept grows with the file. +An exhaustive read therefore costs what a budgeted one costs, and ``row_budget`` is a limit on *work* +rather than the memory guard it used to be. -A partition whose files do not all declare a schema is **materialised**, as before: the schema has -to be inferred from the rows, so the rows have to be kept until it has been. That is line-delimited -formats, and folding them needs accumulators created lazily as columns appear, which is still to do. +What a declared schema buys is not the fold but its sharpness. Parquet footers are read first, so +the columns are known before a row is parsed and each accumulator is chosen up front, and the exact +row count is known too, which is what lets a quality stride be placed over a column not yet seen. + +Without one — line-delimited data — both wait for the data. Columns are created on first sight and +back-filled with the rows they were absent for, and each carries every shape at once until the last +row has gone by and the dtype resolves. That costs a deferred type per column and nothing else; it +does not cost a second pass, and it does not decide from a prefix. The budget is a target rather than a ceiling. :data:`MIN_ROWS_PER_FILE` is the floor every file is read to however thin its share gets, since one sampled below it cannot contribute the columns it @@ -47,7 +50,11 @@ ) from nemo_datasets_plugin.profiler.schema import MAX_COLUMNS, columns_were_capped, derive_features from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits -from nemo_datasets_plugin.profiler.stats import ColumnFold, measure_columns, quote_enumerations +from nemo_datasets_plugin.profiler.stats import ( + ColumnFold, + InferredColumnFold, + quote_enumerations, +) from nemo_platform_plugin.files.dataset_profile import ( ColumnStats, DatasetProfile, @@ -228,59 +235,6 @@ def _unify_schemas(schemas: list[pa.Schema]) -> pa.Schema | None: return None -def _measure( - partition_rows: list[dict], - arrow_schemas: list[pa.Schema], - *, - all_declared: bool, - column_roles: dict[str, str], -) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: - """Derive schema, stats and classification, degrading to structure-only if any of it fails. - - ``all_declared`` says whether *every* file that contributed rows carried a declared schema. When - one did not, the unified schema describes only some of the rows, and using it would erase any - column the schemaless files were the sole witness for — so infer from the rows instead, which - sees all of them. Declared type fidelity (int32 widening to int64) is the cost, and it is the - honest one: a declared schema cannot be asserted over files that declare nothing. - - These stages are pure computation over rows already in memory, so a failure here is either a - profiler bug or data shaped in a way no detector anticipated. Reads are already isolated per - file; leaving this stage unguarded meant one odd value — a chat message whose ``role`` is a number - — could abort an otherwise complete profile from the one place nothing was catching. The - partition's structure (files, splits, row counts) is established by then and stays useful, so the - failure costs its measurements and says so, rather than the entire run. - """ - try: - declared = _unify_schemas(arrow_schemas) if all_declared else None - features = derive_features(partition_rows, declared) - # Statistics and probes together, one column at a time and each isolated. Probes cover every - # column independent of the roles classify is about to assign, so a content signal survives - # a column name the alias table does not know. - measured = measure_columns(features, partition_rows) - # The one probe that reads two columns against each other, so it cannot live on a column's - # accumulator. Folded here over the same rows. - prefix_pair = PrefixPairFold(features) - prefix_pair.update(partition_rows) - classification = classify( - features, - measured.stats, - probes=measured.probes, - prefix_pair=prefix_pair.result(), - column_roles=column_roles, - ) - # Last, because the roles classification assigns are what decide whether a column's values - # may be quoted at all — cardinality only bounds how many. - quote_enumerations(features, measured.stats, measured.vocabularies) - # After classify, so its own reasoning reads first and a column that could not be measured - # is reported as a caveat on the result rather than as part of the case for it. - classification.evidence.extend(_capped_columns_evidence(features)) - classification.evidence.extend(measured.errors) - return features, measured.stats, classification - except Exception as exc: - detail = f"could not measure this partition: {type(exc).__name__}: {exc}" - return [], {}, PartitionClassification(dataset_type="unknown", evidence=[Evidence(kind="error", detail=detail)]) - - @dataclass(frozen=True) class _PartitionOutcome: """One partition plus what it contributes to the dataset-level sampling envelope.""" @@ -348,10 +302,15 @@ class _PartitionFolds: neither. Keeping them side by side is what lets the file loop hand over a batch and forget it. """ - def __init__(self, features: list[FeatureSchema], expected_rows: int | None) -> None: - self.features = features - self._columns = ColumnFold(features, expected_rows) - self._prefix = PrefixPairFold(features) + def __init__(self, features: list[FeatureSchema] | None, expected_rows: int | None) -> None: + # Declared: the columns are known, so the accumulators are chosen now. Inferred: they are + # discovered as they appear and typed once every row has gone by. + self.features = features or [] + self._declared = features is not None + self._columns: ColumnFold | InferredColumnFold = ( + ColumnFold(features, expected_rows) if features is not None else InferredColumnFold(expected_rows) + ) + self._prefix = PrefixPairFold() def update(self, rows: list[dict]) -> None: self._columns.update(rows) @@ -360,13 +319,17 @@ def update(self, rows: list[dict]) -> None: def measure( self, column_roles: dict[str, str] ) -> tuple[list[FeatureSchema], dict[str, ColumnStats], PartitionClassification]: - """Schema, stats and classification from what was folded. Guarded like :func:`_measure`. + """Schema, stats and classification from what was folded. The columns have their own per-column guard inside the fold; this is the wide one, for - anything structural that no single column owns. + anything structural that no single column owns -- a schema that cannot be resolved, + a classifier that trips over a shape no detector anticipated. """ try: - measured = self._columns.finalize() + if isinstance(self._columns, InferredColumnFold): + self.features, measured = self._columns.finalize() + else: + measured = self._columns.finalize() classification = classify( self.features, measured.stats, @@ -400,9 +363,8 @@ def _profile_partition( The reader is resolved per file rather than per partition. Format is a property of a file, and a directory holding two of them is a stray file, not a second dataset; splitting the partition - to keep one scalar ``file_format`` true is what made partition names unstable. Mixed formats - instead flow through to ``_measure``, which infers the schema from rows when not every file - declared one. + to keep one scalar ``file_format`` true is what made partition names unstable. A partition whose + files do not all declare a schema simply infers one, from the rows, as it folds them. An unreadable file (or a format with no registered reader) is isolated: it is named on a :class:`FileError` the envelope collects, contributes no rows, and flips ``scanned_all`` off — it @@ -414,21 +376,15 @@ def _profile_partition( # accumulators must exist before the first batch, and the quality stride must be placed before # the column it strides has been seen. previews = _peek_files(source, entries) - can_fold = all(preview is not None and preview.arrow_schema is not None for preview in previews.values()) - - partition_rows: list[dict] = [] - arrow_schemas: list[pa.Schema] = [ - preview.arrow_schema for preview in previews.values() if preview is not None and preview.arrow_schema - ] - all_declared = True # every file that contributed rows carried a declared schema - folds: _PartitionFolds | None = None - if can_fold: - declared = _unify_schemas(arrow_schemas) - if declared is not None: - folds = _PartitionFolds( - derive_features([], declared), - expected_rows=_expected_rows(previews, _per_file_cap(row_budget, len(entries))), - ) + arrow_schemas = [preview.arrow_schema for preview in previews.values() if preview.arrow_schema is not None] + declared = _unify_schemas(arrow_schemas) if len(arrow_schemas) == len(entries) and arrow_schemas else None + # Declared or not, the partition folds. With a schema the accumulators are chosen up front and + # the exact row count places the quality stride; without one both wait for the data, which costs + # a deferred dtype per column and nothing else. + folds = _PartitionFolds( + derive_features([], declared) if declared is not None else None, + expected_rows=_expected_rows(previews, _per_file_cap(row_budget, len(entries))), + ) rows_scanned = 0 files_read = 0 rows_present: int | None = 0 @@ -450,32 +406,26 @@ def _profile_partition( # Inside the guard: resolving the reader can fail too, and a format with no reader # registered is a file the profiler could not use like any other. reader = get_reader(_format_of(entry.path)) - if folds is not None: - # Streamed: batches are folded and let go, so nothing here grows with the file. - num_rows = previews[entry.path].num_rows # from the footer, read before any row - scanned = 0 - for batch in reader.batches(source, entry, row_cap=row_cap): - folds.update(batch) - scanned += len(batch) - files_read += 1 - rows_scanned += scanned - scanned_all = num_rows is not None and scanned >= num_rows - else: - # Materialised: no declared schema, so the schema has to be inferred from the rows - # and the rows have to be kept until it has been. Phase 4 folds this path too. - result = reader.read(source, entry, row_cap=row_cap) - files_read += 1 - error = result.error - num_rows = result.num_rows - rows_scanned += result.rows_scanned - partition_rows.extend(result.rows) - if result.arrow_schema is None and result.rows: - # Rows with no schema behind them: the unified schema no longer covers the - # partition, so _measure must infer from rows rather than trust a partial one. - all_declared = False - # Exhaustive requires parsing every row; a known footer count alone is not enough, - # and a partial read (corrupt lines skipped) is not exhaustive however many it got. - scanned_all = num_rows is not None and result.rows_scanned >= num_rows and error is None + preview = previews[entry.path] + scanned = 0 + read_errors: list[str] = [] + for batch in reader.batches(source, entry, row_cap=row_cap, errors=read_errors): + folds.update(batch) + scanned += len(batch) + files_read += 1 + rows_scanned += scanned + # A file the reader only partly understood is named, the same as one it could not + # open at all. Folding it silently would make a corrupt shard look complete. + error = "; ".join(read_errors) or None + # A footer knows the count before the read; a line-delimited file only knows it by + # reaching the end, which a capped read does not do. + if preview.num_rows is not None: + num_rows = preview.num_rows + elif row_cap is None or scanned < row_cap: + num_rows = scanned + # Exhaustive requires parsing every row. A known count alone is not enough, and a + # partial read is not exhaustive however many rows it managed to get. + scanned_all = num_rows is not None and scanned >= num_rows and error is None except Exception as exc: # Failure isolation: an unreadable file (or missing reader) keeps its identity, # skips its rows, and does not abort the profile. The reason is recorded rather than @@ -507,12 +457,7 @@ def _profile_partition( num_examples=split_examples if split_counts_known else None, ) ) - if folds is not None: - features, stats, classification = folds.measure(column_roles) - else: - features, stats, classification = _measure( - partition_rows, arrow_schemas, all_declared=all_declared, column_roles=column_roles - ) + features, stats, classification = folds.measure(column_roles) partition = PartitionProfile( name=name, # Observed, not chosen: the partition reports the formats its files turned out to be in diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py index 796bec7728..38b4d3c5aa 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/base.py @@ -63,9 +63,20 @@ def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = No ... def batches( - self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None + self, + source: FileSource, + entry: FileEntry, + *, + row_cap: int | None = None, + errors: list[str] | None = None, ) -> Iterator[list[dict[str, Any]]]: - """The same rows :meth:`read` would return, handed over in chunks and never all at once.""" + """The same rows :meth:`read` would return, handed over in chunks and never all at once. + + ``errors`` collects any reason the read understood less than the whole file, the way + :attr:`ReadResult.error` does for the batched-up path. A generator cannot return one -- + by the time it knows, the caller has consumed everything it yielded -- and the caller has + to know, or a partially parsed file would fold silently and look complete. + """ ... diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py index 8ce2b61672..1ea7da9809 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/jsonl.py @@ -47,14 +47,30 @@ def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: a partition holding one cannot be folded without reading it first.""" return FilePreview() - def batches(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> Iterator[list[dict]]: - """Rows in chunks. Parse errors are silent here -- :meth:`read` is the path that accounts for - them, and the fold path does not reach a file with no declared schema.""" + def batches( + self, + source: FileSource, + entry: FileEntry, + *, + row_cap: int | None = None, + errors: list[str] | None = None, + ) -> Iterator[list[dict]]: + """Rows in chunks, reporting any line it could not read into ``errors``. + + A corrupt line costs that line, never the file -- dropping the whole file would erase its + row count and every column it was the only witness for -- but the caller has to be told, or + a partially parsed file would fold silently and look complete. + """ rows: list[dict] = [] scanned = 0 + unparseable = 0 + first_failure: str | None = None with source.open(entry.path) as stream: - for record, _ in _records(stream): + for record, failure in _records(stream): if record is None: + unparseable += 1 + if first_failure is None: + first_failure = failure continue rows.append(record) scanned += 1 @@ -65,6 +81,8 @@ def batches(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = break if rows: yield rows + if unparseable and errors is not None: + errors.append(f"skipped {unparseable} unparseable line(s); first at {first_failure}") def read(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> ReadResult: rows: list[dict] = [] diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py index 62863bf682..3eed3ccfae 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/readers/parquet.py @@ -26,8 +26,19 @@ def peek(self, source: FileSource, entry: FileEntry) -> FilePreview: parquet_file = pq.ParquetFile(stream) return FilePreview(arrow_schema=parquet_file.schema_arrow, num_rows=parquet_file.metadata.num_rows) - def batches(self, source: FileSource, entry: FileEntry, *, row_cap: int | None = None) -> Iterator[list[dict]]: - """Rows in chunks, so the caller can fold them and let each chunk go.""" + def batches( + self, + source: FileSource, + entry: FileEntry, + *, + row_cap: int | None = None, + errors: list[str] | None = None, + ) -> Iterator[list[dict]]: + """Rows in chunks, so the caller can fold them and let each chunk go. + + ``errors`` is never appended to: parquet either decodes a batch or raises, so there is no + partial understanding to report. + """ scanned = 0 with source.open(entry.path) as stream: parquet_file = pq.ParquetFile(stream) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py index 51c199cd67..3c95f744fb 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/schema.py @@ -143,21 +143,93 @@ def _infer_feature(name: str, values: list[Any]) -> FeatureSchema: return FeatureSchema(name=name, dtype=_scalar_dtype(present)) -def _scalar_dtype(values: list[Any]) -> str: - dtypes: set[str] = set() - for value in values: - if isinstance(value, bool): # bool before int: bool is a subclass of int - dtypes.add("bool") - elif isinstance(value, int): - dtypes.add("int64") - elif isinstance(value, float): - dtypes.add("float64") - elif isinstance(value, str): - dtypes.add("string") - else: - dtypes.add("json") +def _python_dtype(value: Any) -> str: + """The dtype one value implies, on its own.""" + if isinstance(value, bool): # bool before int: bool is a subclass of int + return "bool" + if isinstance(value, int): + return "int64" + if isinstance(value, float): + return "float64" + if isinstance(value, str): + return "string" + return "json" + + +def _resolve_scalar(dtypes: set[str]) -> str: + """The one dtype a column of these observed types has. Ints and floats widen; anything else in + disagreement is ``json``, which is the honest answer for a column that holds two shapes.""" if dtypes <= {"int64", "float64"} and dtypes: return "float64" if "float64" in dtypes else "int64" if len(dtypes) == 1: - return dtypes.pop() + # `next(iter(...))`, never `pop()`: this set belongs to a SchemaFold that is still using it, + # and emptying it made a second call resolve the same column to `json`. + return next(iter(dtypes)) return "json" + + +def _scalar_dtype(values: list[Any]) -> str: + return _resolve_scalar({_python_dtype(value) for value in values}) + + +class SchemaFold: + """One column's schema, folded from values as they arrive rather than decided over all of them. + + The dtype of an inferred column is a whole-column question -- the observed types are unioned and + a disagreement widens to ``json`` -- which is why a partition with no declared schema could not + be folded: an accumulator is chosen *by* dtype, and the dtype is not known until the last row. + + It is a fold, though, and always was. :func:`_infer_feature` is a set union over observed types, + a union over a struct's child keys, and a recursion over a list's flattened elements: state + proportional to the *schema*, not to the row count. This is that computation written incrementally + so a caller can hand over batches and keep none of them. + """ + + def __init__(self, name: str = "") -> None: + self._name = name + self._present = 0 + self._dicts = 0 + self._lists = 0 + self._dtypes: set[str] = set() + self._fields: dict[str, SchemaFold] = {} + self._field_order: list[str] = [] + self._item: SchemaFold | None = None + + def update(self, values: list[Any]) -> None: + for value in values: + if value is None: + continue # a null says nothing about the type; an all-null column resolves to json + self._present += 1 + self._dtypes.add(_python_dtype(value)) + if isinstance(value, dict): + self._dicts += 1 + for key, child in value.items(): + fold = self._fields.get(key) + if fold is None: + fold = SchemaFold(key) + self._fields[key] = fold + self._field_order.append(key) + fold.update([child]) + elif isinstance(value, list): + self._lists += 1 + if self._item is None: + self._item = SchemaFold() + self._item.update(value) + + def finalize(self) -> FeatureSchema: + if not self._present: + return FeatureSchema(name=self._name, dtype="json") + if self._dicts == self._present: + return FeatureSchema( + name=self._name, + dtype="struct", + fields=[self._fields[key].finalize() for key in self._field_order], + ) + if self._lists == self._present: + item = (self._item or SchemaFold()).finalize() + return FeatureSchema( + name=self._name, + dtype="messages" if _is_message_struct(item) else "list", + items=item, + ) + return FeatureSchema(name=self._name, dtype=_resolve_scalar(self._dtypes)) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index 6a108f67f5..a999ad9d59 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -34,6 +34,7 @@ from dataclasses import dataclass from typing import Any +from nemo_datasets_plugin.profiler.schema import MAX_COLUMNS, SchemaFold from nemo_platform_plugin.files.dataset_profile import ( CategoricalStats, ColumnStats, @@ -272,6 +273,16 @@ def _blocks(self) -> dict[str, Any]: """The dtype-specific ``ColumnStats`` blocks. The base column contributes none.""" return {} + def backfill_nulls(self, count: int) -> None: + """Charge this column ``count`` rows in which it was absent. + + A column that first appears in the fiftieth batch was null for every row before it, which is + exactly what a materialising reader computes with ``row.get(name)``. Counted rather than fed + as values, so discovering a column late costs a pair of additions and not a pass. + """ + self.rows += count + self._nulls += count + def vocabulary(self) -> set[Any] | None: """The distinct values, for a column that is a bounded vocabulary. None for one that is not, which is every dtype without a notion of cardinality.""" @@ -343,24 +354,28 @@ def __init__(self, expected_rows: int | None = None) -> None: super().__init__() self._lengths = _LengthHistogram() self._vocabulary = _Vocabulary() - # With the row count known -- parquet footers give it before a row is read -- the quality - # stride can be placed now and each string measured or skipped as it goes by, retaining - # nothing. Without it there is no stride to compute yet, so the strings are held and strided - # at the end. That is the one term here still sized by the column, and it is exactly the - # partitions that have no footer to read. - self._stride = _quality_stride(expected_rows) if expected_rows is not None else None self._quality = _TextQualityCounters() - self._strings: list[str] = [] self._seen = 0 + self._sampled = 0 + # With the row count known -- parquet footers give it before a row is read -- the stride is + # fixed now and the sample is spread evenly over the whole column. Without it there is no + # length to stride over yet, so the stride starts at one and doubles each time the sample + # fills: every row eligible at first, thinning as the column turns out to be long. Each + # sampled row then stands for `stride` rows, which is what keeps the estimate unbiased + # rather than weighted toward the head where sampling was densest. + self._stride = _quality_stride(expected_rows) if expected_rows is not None else 1 + self._adaptive = expected_rows is None def _observe(self, present: list[Any]) -> None: for value in present: if isinstance(value, str): self._lengths.add(len(value)) - if self._stride is None: - self._strings.append(value) - elif self._seen % self._stride == 0: - self._quality.add(value) + if self._seen % self._stride == 0: + self._quality.add(value, self._stride) + self._sampled += 1 + if self._adaptive and self._sampled >= _QUALITY_SAMPLE_ROWS: + self._stride *= 2 + self._sampled = 0 self._seen += 1 self._vocabulary.update(present) @@ -368,7 +383,7 @@ def _blocks(self) -> dict[str, Any]: text = quality = None if self._seen: text = TextStats(chars=self._lengths.quantiles()) - quality = _text_quality(self._strings) if self._stride is None else self._quality.finalize() + quality = self._quality.finalize() return {"text": text, "quality": quality, "categorical": self._vocabulary.finalize()} def vocabulary(self) -> set[Any] | None: @@ -502,6 +517,143 @@ def _accumulator_for(feature: FeatureSchema, expected_rows: int | None = None) - return ColumnAccumulator() +class DeferredAccumulator(ColumnAccumulator): + """A column whose dtype is not known until the last row has gone by. + + An accumulator is normally chosen *by* dtype, which a declared schema gives up front. An inferred + one does not: the observed types are unioned over the whole column and a disagreement widens to + ``json``, so the choice cannot be made while the choosing still matters. Deferring it is the only + resolution that neither reads the data twice nor decides from a prefix and hopes. + + So every shape is measured at once and the answer picked at the end. It costs no more per value + than choosing would have -- a string only ever reaches the string state, an int only the numeric + -- and the state it costs is four bounded structures per column rather than one. A column that + resolves to a shape nothing measured, or to ``json``, simply has no blocks, which is what the + dispatch would have produced for it anyway. + """ + + def __init__(self, name: str, expected_rows: int | None = None) -> None: + super().__init__() + self._schema = SchemaFold(name) + self._string = StringAccumulator(expected_rows) + self._numeric = NumericAccumulator() + self._bool = BoolAccumulator() + self._messages = MessageAccumulator() + + def _observe(self, present: list[Any]) -> None: + self._schema.update(present) + # Routed by python type. Where a dtype resolves to something measurable, every present value + # is of that type by construction -- `_resolve_scalar` only returns `string` when the whole + # column was strings -- so this sees exactly what the chosen accumulator would have seen. + strings = [value for value in present if isinstance(value, str)] + if strings: + self._string._observe(strings) + numbers = [value for value in present if isinstance(value, (int, float)) and not isinstance(value, bool)] + if numbers: + self._numeric._observe(numbers) + bools = [value for value in present if isinstance(value, bool)] + if bools: + self._bool._observe(bools) + lists = [value for value in present if isinstance(value, list)] + if lists: + self._messages._observe(lists) + + def feature(self) -> FeatureSchema: + """The column's schema, as folded.""" + return self._schema.finalize() + + def _blocks(self) -> dict[str, Any]: + dtype = self.feature().dtype + if dtype == "string": + return self._string._blocks() + if dtype == "bool": + return self._bool._blocks() + if dtype == "messages": + return self._messages._blocks() + if _is_numeric(dtype): + return self._numeric._blocks() + return {} + + def vocabulary(self) -> set[Any] | None: + dtype = self.feature().dtype + if dtype == "string": + return self._string.vocabulary() + if dtype == "bool": + return self._bool.vocabulary() + if _is_numeric(dtype): + return self._numeric.vocabulary() + return None + + +class InferredColumnFold: + """A partition's columns, discovered as they appear and typed once they have all gone by. + + The counterpart to :class:`ColumnFold` for data that declares no schema. Columns are created on + first sight and back-filled with the rows they were absent for, which is what makes the result + identical to inferring the schema first and measuring second -- a row without the key genuinely + holds a null for it. + """ + + def __init__(self, expected_rows: int | None = None) -> None: + self._expected_rows = expected_rows + self._accumulators: dict[str, DeferredAccumulator] = {} + self._order: list[str] = [] + self._failed: dict[str, Evidence] = {} + self._rows_seen = 0 + + def update(self, rows: list[dict[str, Any]]) -> None: + for row in rows: + for name in row: + if name in self._accumulators or len(self._accumulators) >= MAX_COLUMNS: + continue + accumulator = DeferredAccumulator(name, self._expected_rows) + accumulator.backfill_nulls(self._rows_seen) + self._accumulators[name] = accumulator + self._order.append(name) + for name in self._order: + if name in self._failed: + continue + try: + self._accumulators[name].update([row.get(name) for row in rows]) + except Exception as exc: + self._failed[name] = Evidence( + kind="error", + detail=f"column {name!r} could not be measured: {type(exc).__name__}: {exc}", + ) + self._rows_seen += len(rows) + + def finalize(self) -> tuple[list[FeatureSchema], ColumnMeasurements]: + features: list[FeatureSchema] = [] + stats: dict[str, ColumnStats] = {} + probes: dict[str, ColumnProbes] = {} + vocabularies: dict[str, set[Any]] = {} + errors: list[Evidence] = [] + for name in self._order: + failure = self._failed.get(name) + if failure is not None: + errors.append(failure) + continue + accumulator = self._accumulators[name] + try: + features.append(accumulator.feature()) + column, probe = accumulator.finalize() + except Exception as exc: + errors.append( + Evidence( + kind="error", + detail=f"column {name!r} could not be summarised: {type(exc).__name__}: {exc}", + ) + ) + continue + probes[name] = probe + vocabulary = accumulator.vocabulary() + if vocabulary is not None: + vocabularies[name] = vocabulary + if column is not None: + stats[name] = column + return features, ColumnMeasurements(stats=stats, probes=probes, vocabularies=vocabularies, errors=errors) + + def _is_numeric(dtype: str) -> bool: return dtype.startswith(("int", "uint", "float")) @@ -608,28 +760,6 @@ def _at(self, percentile: int) -> int: _QUALITY_SAMPLE_ROWS = 50_000 -def _quality_sample(strings: list[str]) -> list[str]: - """The rows to measure quality over: all of them, or an evenly strided subset. - - Strided rather than random, because two runs over the same bytes must agree. Randomness is what - ``SamplingInfo.seed`` existed to make reproducible, and that field was deleted on the grounds - that the profiler makes no random choices and a seed would be theatre -- which should stay true. - - Strided rather than the head, because shards arrive sorted often enough that the first rows of a - large column are not a sample of it. A stride costs the same, needs no state, and spreads the - sample across the whole column. - - A stride can in principle alias against periodic data: a set with two rows per prompt, sampled - at stride two, sees one phase of every pair. Measured on exactly that shape -- HelpSteer2 rates - two responses per prompt -- the two phases agree to 0.35% on `whitespace_ratio` and differ by at - most 8% on the other two, whose values there are 0.0003 and 0.0025. That is well inside the band - these estimates already carry near zero, and does not buy a block-sampling scheme to avoid. - """ - if len(strings) <= _QUALITY_SAMPLE_ROWS: - return strings - return strings[:: len(strings) // _QUALITY_SAMPLE_ROWS] - - def _whitespace_count(text: str) -> int: """Whitespace characters, matching ``\\s`` exactly. @@ -670,12 +800,18 @@ def __init__(self) -> None: self._repetition = 0.0 self._rows = 0 - def add(self, text: str) -> None: - self._chars += len(text) - self._whitespace += _whitespace_count(text) - self._non_ascii += _non_ascii_count(text) - self._repetition += _repetition_score(text) - self._rows += 1 + def add(self, text: str, weight: int = 1) -> None: + """Fold one sampled row in, standing for ``weight`` rows of the column. + + The weight is what keeps a *varying* sample rate honest. Sampling one row in four and + counting it once would let a densely sampled stretch outvote a sparsely sampled one; counting + it four times estimates the population sums instead, and the ratios come out unbiased. + """ + self._chars += weight * len(text) + self._whitespace += weight * _whitespace_count(text) + self._non_ascii += weight * _non_ascii_count(text) + self._repetition += weight * _repetition_score(text) + self._rows += weight def finalize(self) -> TextQuality: return TextQuality( @@ -690,26 +826,6 @@ def _quality_stride(rows: int) -> int: return max(1, rows // _QUALITY_SAMPLE_ROWS) -def _text_quality(strings: list[str]) -> TextQuality: - sample = _quality_sample(strings) - total_chars = 0 - whitespace = 0 - non_ascii = 0 - repetition_sum = 0.0 - for value in sample: - total_chars += len(value) - whitespace += _whitespace_count(value) - non_ascii += _non_ascii_count(value) - repetition_sum += _repetition_score(value) - return TextQuality( - whitespace_ratio=whitespace / total_chars if total_chars else 0.0, - non_ascii_ratio=non_ascii / total_chars if total_chars else 0.0, - # Every denominator is the sample's own, never the column's: each ratio is an estimate over - # the rows that were actually scanned, which is what keeps it unbiased rather than diluted. - repetition_score=repetition_sum / len(sample) if sample else 0.0, - ) - - def _count_matches(pattern: re.Pattern[str], text: str) -> int: return sum(1 for _ in pattern.finditer(text)) diff --git a/plugins/nemo-datasets/tests/test_classify.py b/plugins/nemo-datasets/tests/test_classify.py index ef1e72a533..304852b5d8 100644 --- a/plugins/nemo-datasets/tests/test_classify.py +++ b/plugins/nemo-datasets/tests/test_classify.py @@ -20,7 +20,7 @@ def _probes(features, rows): def classify_rows(features, stats, rows, **kwargs): """Classify from rows the way the pipeline does: fold first, then interpret the folds.""" - prefix = PrefixPairFold(features) + prefix = PrefixPairFold() prefix.update(rows) return classify(features, stats, probes=_probes(features, rows), prefix_pair=prefix.result(), **kwargs) diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index b2911cfa8d..14a6189af3 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -12,7 +12,7 @@ import pytest from nemo_datasets_plugin.profiler.file_source import FileEntry, LocalFileSource from nemo_datasets_plugin.profiler.partition import group_partitions -from nemo_datasets_plugin.profiler.pipeline import _expected_rows, _measure, _peek_files, profile +from nemo_datasets_plugin.profiler.pipeline import _expected_rows, _peek_files, profile from nemo_datasets_plugin.profiler.readers.base import FilePreview from nemo_datasets_plugin.profiler.splits import infer_data_files, resolve_splits from nemo_platform_plugin.files.dataset_profile import DatasetProfile @@ -777,20 +777,28 @@ def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): assert result.partitions[0].stats_complete is False -def test_measure_infers_from_rows_when_some_files_declared_no_schema(): - # derive_features uses the declared schema *if present at all*, so a group where only some files - # declare one erased every column the schemaless files were the sole witness for. That is the - # defect `_split_by_format` worked around by making partitions format-homogeneous; the fix - # belongs in schema derivation, and dropping that homogeneity is what makes this path reachable. - declared = pa.schema([pa.field("prompt", pa.string())]) - rows = [{"prompt": "a"}, {"prompt": "b", "extra": "only in the schemaless file"}] +def test_a_column_only_a_schemaless_file_witnessed_survives(tmp_path): + # A group where only some files declare a schema used to trust that schema and erase every column + # the schemaless files were the sole witness for. Now the partition infers its schema from the + # rows as it folds them, so the sole witness is heard. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a"}]) + (tmp_path / "train-00001-of-00002.jsonl").write_text(json.dumps({"prompt": "b", "extra": "only here"})) - features, stats, _ = _measure(rows, [declared], all_declared=False, column_roles={}) - assert [f.name for f in features] == ["prompt", "extra"] # the sole witness survives - assert set(stats) <= {f.name for f in features} + part = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] + + assert [f.name for f in part.features] == ["prompt", "extra"] + assert set(part.stats) <= {f.name for f in part.features} + + +def test_a_declared_schema_is_trusted_when_it_covers_every_file(tmp_path): + # The other half: when every file declares one, the schema is authoritative and the rows are not + # consulted for it. That is what lets the partition fold with its accumulators chosen up front. + _write_parquet(tmp_path / "train-00000-of-00002.parquet", [{"prompt": "a"}]) + _write_parquet(tmp_path / "train-00001-of-00002.parquet", [{"prompt": "b"}]) + + part = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions[0] - features, _, _ = _measure(rows, [declared], all_declared=True, column_roles={}) - assert [f.name for f in features] == ["prompt"] # declared schema trusted when it covers everything + assert [f.name for f in part.features] == ["prompt"] def test_stats_completeness_is_per_partition(tmp_path): diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 70d567ae28..70a8fe105c 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -12,10 +12,8 @@ _MAX_VOCABULARY_VALUE_CHARS, _MAX_VOCABULARY_VALUES, _NON_ASCII_RUN, - _QUALITY_SAMPLE_ROWS, _WHITESPACE_RUN, _non_ascii_count, - _quality_sample, _whitespace_count, measure_columns, quote_enumerations, @@ -259,16 +257,35 @@ def test_quality_fast_paths_are_the_same_measurement_as_the_regexes(text): assert _non_ascii_count(text) == sum(1 for _ in _NON_ASCII_RUN.finditer(text)) -def test_quality_sample_is_bounded_strided_and_deterministic(): - under = [f"r{i}" for i in range(_QUALITY_SAMPLE_ROWS)] - assert _quality_sample(under) is under # nothing to sample: measured in full +def test_a_known_row_count_strides_evenly_and_deterministically(monkeypatch): + # With the row count known up front the stride is fixed, so the sample is spread evenly over the + # whole column -- and two runs over the same bytes agree, which is why no RNG is involved. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 10) + values = ["clean text"] * 100 + ["aaaaaaaaaaaa"] * 100 - over = [f"r{i}" for i in range(_QUALITY_SAMPLE_ROWS * 3)] - sample = _quality_sample(over) - assert len(sample) <= _QUALITY_SAMPLE_ROWS + 1 - assert _quality_sample(over) == sample # no RNG, so no seed to record and no run-to-run drift - # Spans the column rather than its head: the last sampled row is near the end. - assert over.index(sample[-1]) >= len(over) - 3 + def quality(expected_rows): + acc = stats_module.StringAccumulator(expected_rows) + acc.update(values) + return acc.finalize()[0].quality.repetition_score + + assert quality(len(values)) == quality(len(values)) # deterministic + # Half the column is corrupt and the stride spans it, so the estimate lands near a half. + assert 0.4 <= quality(len(values)) <= 0.6 + + +def test_an_unknown_row_count_thins_as_it_goes_and_stays_unbiased(monkeypatch): + # No footer, so no length to stride over: the stride starts at one and doubles as the sample + # fills. Sampling is then densest at the head, which would skew the answer -- weighting each + # sampled row by the stride it stood for is what corrects it. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 10) + values = ["clean text"] * 500 + ["aaaaaaaaaaaa"] * 500 + + acc = stats_module.StringAccumulator(None) + acc.update(values) + score = acc.finalize()[0].quality.repetition_score + + assert acc._stride > 1 # it did thin + assert 0.35 <= score <= 0.65 # ...and still found roughly half the column corrupt def test_quality_is_measured_across_the_column_not_its_head(monkeypatch): From 6b94f31bd527dcf2aa4d43fa5fae6a949a4a764b Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 7 Aug 2026 21:08:57 -0400 Subject: [PATCH 43/44] feat(datasets): read everything by default, and say what completeness means MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 5, the last of the streaming spec. The engine can finally honour the contract change, so it lands. `DEFAULT_ROW_BUDGET` is now None. The budget existed to keep a materialised partition off the heap, and nothing is materialised: a fold's memory is flat in rows, so an exhaustive read costs what a short one costs. The default should not answer the question worse than it can be answered. The demo now reads all 67,551 rows at 11.5 MB peak, both partitions `rows_complete`. `stats_complete` becomes `rows_complete`, which is what it measured all along. The old name promised more than it delivered -- `Quantiles` and `TextQuality` are estimates by construction however much was read, and each says so itself. Whether a number is exact is a property of that number; this says only whether anything was missed on the way in. `SamplingInfo.row_budget` is dropped. It is an input, not a finding, and the finding is already there: `rows_scanned` against `rows_present` says a read was short, and the only other cause -- a file that failed -- is named on `file_errors`. The parameter survives on `profile()` and in the step config for a caller who wants a shorter run. Two deviations from the spec, both because the design moved under it. `rows_measured` is not added. §11 conceived it as the denominator behind an estimate's confidence, sized `min(reservoir, rows_scanned)`. There is no reservoir: quantiles read off a histogram over every scanned row, and the quality sample is a per-column stride. There is no single dataset-level number left for the field to hold, and inventing one would be worse than the absence. `_per_file_cap` and `MIN_ROWS_PER_FILE` survive, which Q3 guessed they might not. They look like budget-splitting machinery invented for the memory problem, but dividing a budget across files was always about *coverage*: reading files in order until a total ran out leaves the later ones unopened, which hides the columns only they witness. Same hole, different route. Memory was never what that arithmetic was protecting. Correction to the spec's runtime projection while I am here. §2 put an exhaustive OpenMathReasoning at ~3 min on the strength of a 450 M chars/s "cheap path", which was measured on a stripped loop -- no probes, no vocabulary, no histogram. The real fold runs at 13 M chars/s below the quality stride's threshold and 22 M above it, which puts that dataset nearer 50 minutes. Memory was the goal and memory is flat; runtime scaling with the dataset was accepted going in. But the number in the spec was wrong and is now the measured one. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 52 ++++++------- .../tests/files/test_dataset_profile.py | 14 ++-- .../nemo_datasets_plugin/profiler/pipeline.py | 47 ++++++------ .../nemo_datasets_plugin/tasks/profile/run.py | 15 ++-- plugins/nemo-datasets/tests/test_pipeline.py | 75 ++++++++++++------- .../nemo-datasets/tests/test_profile_task.py | 13 ++-- 6 files changed, 118 insertions(+), 98 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 295445c854..581ef74c48 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -304,7 +304,7 @@ class CategoricalStats(BaseModel): "vocabulary (label | provenance | meta | rank) and it holds at most 32 of them. Cardinality " "alone cannot be the gate: it inverts on small data, where every column holds few distinct " "values — free text included — so a three-row dataset had its prompts stored verbatim. A role " - "says what a column *is*, at any size. Read `PartitionProfile.stats_complete` to know whether " + "says what a column *is*, at any size. Read `PartitionProfile.rows_complete` to know whether " "this is the whole vocabulary or only what the sampled rows showed." ), ) @@ -469,17 +469,19 @@ class PartitionProfile(BaseModel): "omitted); keys are a subset of the top-level `features` names." ), ) - stats_complete: bool = Field( + rows_complete: bool = Field( description=( - "True => `features`, `stats` and `classification` were computed over every row of every " - "file in THIS partition: proven facts, not estimates. Only then can a consumer assert " - "enum / required in a bridged JSON Schema, or read a verifiability coverage of 1.0 as " - "literal. It speaks to *rows read*, not to every number being exact: `TextQuality` and " - "`Quantiles` are estimates by construction however much was read, each bounded for the " - "cost reasons its own docstring gives, and `Quantiles.max` is exact regardless. " - "Scoped to the partition because that is where it is decided — a corrupt shard " - "in one partition says nothing about the measurements in another, and a fileset-wide " - "flag quietly downgraded every partition to the worst one." + "True => every row of every file in THIS partition was read. Only then can a consumer " + "assert enum / required in a bridged JSON Schema, or read a verifiability coverage of " + "1.0 as literal.\n\n" + "Named for what it measures. It was `stats_complete`, which promised more than it " + "delivered: `Quantiles` and `TextQuality` are estimates by construction however much was " + "read, each bounded for the cost reasons its own docstring gives. Whether a number is " + "exact is a property of that number, and every one of them says so; this says only " + "whether anything was missed on the way in.\n\n" + "Scoped to the partition because that is where it is decided — a corrupt shard in one " + "partition says nothing about the measurements in another, and a fileset-wide flag " + "quietly downgraded every partition to the worst one." ), ) classification: PartitionClassification @@ -501,15 +503,20 @@ class SamplingInfo(BaseModel): """How much of the data the profile is based on — coverage, stated as numbers. Deliberately carries no ``exhaustive`` flag. That bit was answering two questions at once: "are - these measurements facts or estimates?", which is a property of a *partition* and now lives on - ``PartitionProfile.stats_complete``, and "did I see all the data?", which is this block's job and - needs numerators and denominators rather than a boolean. It also folded together causes that - call for different people to act — a row cap is the caller's choice, a corrupt shard is the data + these measurements facts or estimates?", which is a property of each measurement and is now + stated by each of them, and "did I see all the data?", which is this block's job and needs + numerators and denominators rather than a boolean. It also folded together causes that call for + different people to act — a short read is the caller's choice, a corrupt shard is the data owner's problem, and a missing reader is ours. - The dataset-wide question is still one expression away, and now says which half failed:: + Nor does it record the caller's row limit. Reading everything is now the default and costs what + reading some of it costs, so a short read is unusual — and when it happens ``rows_scanned`` + against ``rows_present`` already says so. *Why* is not the profile's business: a limit is an + input, and the only other cause is a file that failed, which is named on ``file_errors``. - all(p.stats_complete for p in profile.partitions) and not profile.file_errors + The dataset-wide question is still one expression away, and still says which half failed:: + + all(p.rows_complete for p in profile.partitions) and not profile.file_errors """ rows_scanned: int = Field(description="Total rows actually parsed across all files.") @@ -548,17 +555,6 @@ class SamplingInfo(BaseModel): "stops being derivable the moment coverage is partial, which is the only time it is read." ), ) - row_budget: int | None = Field( - default=None, - description=( - "Rows the caller allowed per partition, None for an unbounded read. A budget rather than a " - "per-file cap because the cost is per partition: a per-file cap made peak memory scale with " - "shard count, so the same dataset resharded from 100 files to 10,000 went from megabytes to " - "gigabytes without holding any more data. Not a hard ceiling: every file is still read at " - "least a few rows, since one sampled too thinly cannot contribute the columns it alone " - "witnesses, so a partition with very many files may exceed its budget." - ), - ) class DatasetProfile(BaseModel): diff --git a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py index 9bf0efcf59..ed2e8359c0 100644 --- a/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py +++ b/packages/nemo_platform_plugin/tests/files/test_dataset_profile.py @@ -36,11 +36,11 @@ created_at: 2026-07-08T22:05:12Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 2112, rows_present: 3201061, - files_read: 33, files_present: 33, bytes_present: 31821490182, row_budget: 4096} + files_read: 33, files_present: 33, bytes_present: 31821490182} partitions: - name: "" file_formats: [parquet] - stats_complete: false + rows_complete: false splits: - {name: train, canonical: train, num_examples: 3200861, num_files: 32, size_bytes: 31819412254, data_files: 'train*.parquet'} @@ -77,11 +77,11 @@ created_at: 2026-07-08T22:41:37Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 1024, rows_present: 46189, - files_read: 2, files_present: 2, bytes_present: 27055195, row_budget: 1024} + files_read: 2, files_present: 2, bytes_present: 27055195} partitions: - name: "" file_formats: [parquet] - stats_complete: false + rows_complete: false splits: - {name: train, canonical: train, num_examples: 43835, num_files: 1, size_bytes: 25670988, data_files: 'train*.parquet'} @@ -118,11 +118,11 @@ created_at: 2026-07-09T10:12:45Z profiler_info: {name: nemo-dataset-profiler, version: 0.1.0} sampling: {rows_scanned: 1024, rows_present: 21362, - files_read: 2, files_present: 2, bytes_present: 19459677, row_budget: 1024} + files_read: 2, files_present: 2, bytes_present: 19459677} partitions: - name: "" file_formats: [parquet] - stats_complete: false + rows_complete: false splits: - {name: train, canonical: train, num_examples: 20324, num_files: 1, size_bytes: 18495985, data_files: 'train*.parquet'} @@ -180,7 +180,7 @@ def _build_profile() -> DatasetProfile: partitions=[ PartitionProfile( file_formats=["parquet"], - stats_complete=False, + rows_complete=False, splits=[ SplitProfile( name="train", diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 4742f7fcc6..4199c2867f 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -13,8 +13,8 @@ shards. Every partition is **folded**: batches are measured and let go, and nothing kept grows with the file. -An exhaustive read therefore costs what a budgeted one costs, and ``row_budget`` is a limit on *work* -rather than the memory guard it used to be. +An exhaustive read therefore costs what a short one costs, which is why reading everything is the +default. ``row_budget`` survives only as a way to ask for a shorter run. What a declared schema buys is not the fold but its sharpness. Parquet footers are read first, so the columns are known before a row is parsed and each accumulator is chosen up front, and the exact @@ -25,11 +25,11 @@ row has gone by and the dtype resolves. That costs a deferred type per column and nothing else; it does not cost a second pass, and it does not decide from a prefix. -The budget is a target rather than a ceiling. :data:`MIN_ROWS_PER_FILE` is the floor every file is -read to however thin its share gets, since one sampled below it cannot contribute the columns it -alone witnesses. Files smaller than their share are read to the end and keep exact row counts, so a -budgeted profile of a small dataset is still complete. Pass ``row_budget=None`` for a genuinely -exhaustive scan. +A caller who does ask for one gets a target rather than a ceiling. :data:`MIN_ROWS_PER_FILE` is the +floor every file is read to however thin its share gets, since one sampled below it cannot +contribute the columns it alone witnesses. That division outlived the memory problem it was invented +for: reading files in order until a total ran out would leave the later ones unopened, which is the +same coverage hole by another route. """ from __future__ import annotations @@ -70,18 +70,17 @@ PROFILER_NAME = "nemo-dataset-profiler" PROFILER_VERSION = "0.1.0" -# Rows a partition may read, in total, by default. Every file is still opened — head-sampling a -# *subset of files* would hide columns that appear only in later shards — but the budget is divided -# across them, so peak memory tracks the budget rather than the shard count. A per-file cap put the -# knob on the wrong axis: at 1000 rows each, resharding a dataset from 100 files to 10,000 took peak -# heap from 135 MB to 13.5 GB while describing exactly the same data. Ten thousand rows is ample for -# the statistics computed here (length quantiles, rates, cardinality); pass ``row_budget=None`` for a -# genuinely exhaustive scan. -DEFAULT_ROW_BUDGET = 10_000 - -# Rows read from a file however thin the budget gets. Below this a file cannot contribute the columns -# it alone witnesses, which is the whole reason every file is opened rather than a subset sampled. It -# is what makes the budget a target rather than a ceiling: 10,000 shards read this many each. +# Read everything. The budget existed to keep a materialised partition off the heap, and nothing is +# materialised any longer -- a fold's memory is flat in rows, so an exhaustive read costs what a +# short one costs. What it bounded was never really rows, it was risk. +DEFAULT_ROW_BUDGET = None + +# Rows read from a file however thin a caller-supplied budget gets. Below this a file cannot +# contribute the columns it alone witnesses, which is the whole reason every file is opened rather +# than a subset sampled. It is what makes a budget a target rather than a ceiling: 10,000 shards read +# this many each. It survives the default going unbounded because it never had anything to do with +# memory -- dividing a budget across files is about *coverage*, and reading files in order until a +# total ran out would leave the later ones unopened. MIN_ROWS_PER_FILE = 10 @@ -103,9 +102,10 @@ def profile( ) -> DatasetProfile: """Profile the dataset behind ``source`` into a ``DatasetProfile``. - ``row_budget`` bounds how many rows each *partition* reads in total, divided across its files; - ``None`` reads every row, which is exact but scales memory with the dataset. Files smaller than - their share are read to the end, so a budgeted profile of a small dataset stays complete. + ``row_budget`` bounds how many rows each *partition* reads in total, divided across its files. + It defaults to ``None``, which reads every row: memory is flat in rows either way, so the only + thing a budget buys now is a shorter run. Files smaller than their share are read to the end, so + a budgeted profile of a small dataset is still complete. ``column_roles`` maps a column name to a role the caller is asserting, for datasets whose column names the role table does not recognize. Hints take precedence over name detection but still have @@ -173,7 +173,6 @@ def profile( # Summing the splits would miss the unreadable files, which never reach a partition. bytes_present=sum(entry.size_bytes for entry in data_entries) + sum(entry.size_bytes for entry in unreadable_entries), - row_budget=row_budget, ) return DatasetProfile( created_at=created_at, @@ -468,7 +467,7 @@ def _profile_partition( stats=stats, # Scoped to this partition, which is where it was decided all along: `partition_scanned` is # the value that already gated whether `categorical.values` could quote a proven enumeration. - stats_complete=partition_scanned, + rows_complete=partition_scanned, classification=classification, ) return _PartitionOutcome( diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py index 472091e179..6c977d4751 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/tasks/profile/run.py @@ -110,13 +110,14 @@ def _build_source(config: dict) -> FileSource: def _resolve_row_budget(config: dict) -> int | None: """Rows the profiler may read per partition, from the step config. - Defaults to the profiler's budget rather than an exhaustive read: uncapped, a partition holds - every row of every file in memory at roughly 20x the on-disk parquet size, which is what makes a - large fileset kill the job outright. A budgeted profile keeps exact row counts from the parquet - footers and reports ``stats_complete: false`` for the measurements, which is the trade the - sampling contract exists to describe. - - ``0`` asks for every row; use it when a proven value enumeration matters more than the cost. + Defaults to reading everything, which it did not used to. A partition was materialised before it + was measured, so an uncapped run held every row of every file at roughly 20x the on-disk parquet + size and a large fileset killed the job outright. The profiler folds now: memory is flat in rows, + so the only thing a budget buys is a shorter run, and the default should not be to answer a + question worse than it can be answered. + + ``0`` and ``null`` both ask for every row -- the same thing the default does -- and are kept so a + caller that was setting them explicitly still means what it meant. """ if "row_budget" not in config: return DEFAULT_ROW_BUDGET diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 14a6189af3..330928eb5c 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -205,7 +205,7 @@ def test_the_folded_and_materialised_paths_measure_the_same_thing(tmp_path): def test_an_exhaustive_fold_does_not_cost_more_than_a_budgeted_one(tmp_path): # The point of the whole exercise: reading every row costs what reading some of them costs, so - # the budget stops being a memory guard. Same measurements, and `stats_complete` finally true. + # the budget stops being a memory guard. Same measurements, and `rows_complete` finally true. _write_parquet(tmp_path / "train.parquet", [{"t": f"row {i}" * (i % 5 + 1)} for i in range(5000)]) budgeted = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=500) @@ -213,8 +213,8 @@ def test_an_exhaustive_fold_does_not_cost_more_than_a_budgeted_one(tmp_path): assert budgeted.sampling.rows_scanned == 500 assert exhaustive.sampling.rows_scanned == 5000 - assert budgeted.partitions[0].stats_complete is False - assert exhaustive.partitions[0].stats_complete is True + assert budgeted.partitions[0].rows_complete is False + assert exhaustive.partitions[0].rows_complete is True # Exact where it claims to be exact: the longest row is found by reading all of them. assert exhaustive.partitions[0].stats["t"].text.chars.max >= budgeted.partitions[0].stats["t"].text.chars.max @@ -293,8 +293,8 @@ def test_profile_parquet_dataset_builds_envelope(tmp_path): # A budgeted run over files that all fit under their share is still a complete scan, which is # why the budget and the outcome are separate fields. - assert partition.stats_complete is True - assert result.sampling.row_budget == 10_000 + assert partition.rows_complete is True + assert result.sampling.rows_scanned == result.sampling.rows_present # exhaustive by default assert result.sampling.rows_scanned == 3 assert result.sampling.rows_present == 3 assert result.sampling.files_read == result.sampling.files_present == 2 @@ -368,7 +368,7 @@ def test_profile_keeps_a_mixed_format_directory_as_one_partition(tmp_path): # `question`, which only the schemaless file witnesses -- the defect the split worked around. assert sorted(f.name for f in partition.features) == ["prompt", "question"] assert result.sampling.rows_scanned == 3 # 1 parquet + 2 jsonl, each counted once - assert partition.stats_complete is True + assert partition.rows_complete is True def test_root_files_and_a_directory_named_default_stay_distinct(): @@ -448,7 +448,7 @@ def test_profile_isolates_unreadable_files(tmp_path): assert splits["test"].num_examples is None # unreadable -> count unknown, not a crash assert [e.path for e in result.file_errors] == ["test-00000-of-00001.parquet"] # named, with a reason assert result.file_errors[0].error - assert result.partitions[0].stats_complete is False # a file could not be fully parsed + assert result.partitions[0].rows_complete is False # a file could not be fully parsed assert result.sampling.rows_present is None assert result.sampling.files_read == 1 # one file was actually read; the other never opened assert result.sampling.files_present == 2 # ...out of two that were there to read @@ -460,8 +460,7 @@ def test_profile_row_budget_bounds_reads_and_says_so(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=4) assert result.sampling.rows_scanned == 4 - assert result.sampling.row_budget == 4 - assert result.partitions[0].stats_complete is False # 4 of 10 rows is not a full scan + assert result.partitions[0].rows_complete is False # 4 of 10 rows is not a full scan # The footer knows the total even though the cap stopped the read. Gating this on completeness # nulled it exactly when it carried information: "4 of 10" is a ratio, "4 of unknown" is not. assert result.sampling.rows_present == 10 @@ -473,8 +472,7 @@ def test_profile_uncapped_read_is_a_full_scan(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=None) - assert result.sampling.row_budget is None - assert result.partitions[0].stats_complete is True + assert result.partitions[0].rows_complete is True assert result.sampling.rows_scanned == result.sampling.rows_present == 10 @@ -486,7 +484,7 @@ def test_profile_cap_larger_than_a_jsonl_file_keeps_it_exhaustive(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=1000) assert result.partitions[0].splits[0].num_examples == 2 - assert result.partitions[0].stats_complete is True + assert result.partitions[0].rows_complete is True assert result.sampling.rows_present == 2 @@ -520,7 +518,7 @@ def test_profile_ignores_non_data_files_without_penalty(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert result.partitions[0].stats_complete is True + assert result.partitions[0].rows_complete is True assert result.file_errors == [] assert result.sampling.files_present == 1 # the README and LICENSE are not data, counted nowhere # Nor does their weight land on the dataset: a card is not part of what has to be moved. @@ -554,7 +552,7 @@ def test_profile_records_a_partial_jsonl_read(tmp_path): assert result.partitions[0].splits[0].num_examples == 2 # the readable rows survived assert [e.path for e in result.file_errors] == ["train.jsonl"] assert "line 2" in result.file_errors[0].error - assert result.partitions[0].stats_complete is False # a line was lost, so not a full scan + assert result.partitions[0].rows_complete is False # a line was lost, so not a full scan def test_profile_classifies_roles_type_and_verifiability(tmp_path): @@ -628,9 +626,9 @@ def test_a_measurement_failure_does_not_look_like_a_read_failure(tmp_path, monke assert result.file_errors == [] # the file was fine; the data was odd assert [e.kind for e in result.partitions[0].classification.evidence] == ["error"] - # `stats_complete` speaks to rows read, and every row *was* read -- so it stays True even though + # `rows_complete` speaks to rows read, and every row *was* read -- so it stays True even though # there are no stats. Pinned as it stands; the field means what it says once Phase 5 renames it. - assert result.partitions[0].stats_complete is True + assert result.partitions[0].rows_complete is True def test_one_unmeasurable_column_does_not_cost_the_partition_its_classification(tmp_path, monkeypatch): @@ -684,6 +682,32 @@ def poison_one_partition(features, stats, **kwargs): assert partitions["good"].stats # a neighbour's bad data costs this partition nothing +def test_reading_everything_is_the_default(tmp_path): + # The point of the whole exercise. The budget existed to keep a materialised partition off the + # heap; nothing is materialised, so the default should not answer the question worse than it can + # be answered. + _write_parquet(tmp_path / "train.parquet", [{"t": f"row {i}"} for i in range(5_000)]) + + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.sampling.rows_scanned == 5_000 == result.sampling.rows_present + assert result.partitions[0].rows_complete is True + + +def test_rows_complete_speaks_to_rows_read_not_to_exactness(tmp_path): + # It was `stats_complete`, which promised more than it delivered: quantiles and quality ratios + # are estimates by construction, whatever it says. Renamed to what it actually measures. + _write_parquet(tmp_path / "train.parquet", [{"t": f"row {i}"} for i in range(100)]) + + short = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=10) + whole = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert short.partitions[0].rows_complete is False # ten of a hundred rows + assert whole.partitions[0].rows_complete is True + # True either way, and it is the measurements themselves that say whether they are exact. + assert whole.partitions[0].stats["t"].text.chars.max == max(len(f"row {i}") for i in range(100)) + + def test_profile_is_deterministic(tmp_path): _write_parquet(tmp_path / "train-00000-of-00001.parquet", [{"a": 1}, {"a": 2}]) source = LocalFileSource(tmp_path) @@ -699,7 +723,7 @@ def test_profile_tolerates_non_object_jsonl_lines(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) assert result.partitions[0].splits[0].num_examples == 2 # objects counted, stray array dropped - assert result.partitions[0].stats_complete is True + assert result.partitions[0].rows_complete is True def test_profile_survives_a_hostile_directory(tmp_path): @@ -722,7 +746,7 @@ def test_profile_survives_a_hostile_directory(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise # Nothing here is exhaustive, and the profile says so rather than looking clean. - assert result.partitions[0].stats_complete is False + assert result.partitions[0].rows_complete is False assert result.sampling.rows_present is None # One channel for every file the profiler could not use, whether or not a partition grouped it: # the .csv it never read, the corrupt shard, and the jsonl it only partly parsed. @@ -774,7 +798,7 @@ def test_profile_isolates_detected_format_with_no_reader(tmp_path, monkeypatch): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) # must not raise assert "extra.xyz" in {e.path for e in result.file_errors} # named, not silently dropped - assert result.partitions[0].stats_complete is False + assert result.partitions[0].rows_complete is False def test_a_column_only_a_schemaless_file_witnessed_survives(tmp_path): @@ -801,7 +825,7 @@ def test_a_declared_schema_is_trusted_when_it_covers_every_file(tmp_path): assert [f.name for f in part.features] == ["prompt"] -def test_stats_completeness_is_per_partition(tmp_path): +def test_rows_completeness_is_per_partition(tmp_path): # A corrupt shard in one partition says nothing about the measurements in another, but a # fileset-wide flag downgraded every partition to the worst one. It was never even the value # that gated quoting a proven enumeration -- that was decided per partition and never stored. @@ -812,10 +836,10 @@ def test_stats_completeness_is_per_partition(tmp_path): partitions = {p.name: p for p in profile(LocalFileSource(tmp_path), created_at=FIXED_TIME).partitions} - assert partitions["main"].stats_complete is True - assert partitions["socratic"].stats_complete is False + assert partitions["main"].rows_complete is True + assert partitions["socratic"].rows_complete is False # Quoting is decided by role, not by completeness, so both keep their label vocabulary -- - # stats_complete is what tells a consumer whether socratic's list is the whole of it. + # rows_complete is what tells a consumer whether socratic's list is the whole of it. assert partitions["main"].stats["label"].categorical.values == ["False", "True"] assert partitions["socratic"].stats["label"].categorical.values == ["False", "True"] @@ -826,11 +850,11 @@ def test_dataset_wide_completeness_is_one_expression(tmp_path): # says *which* half failed, which the single bit could not. _write_parquet(tmp_path / "train.parquet", [{"a": 1}]) clean = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert all(p.stats_complete for p in clean.partitions) and not clean.file_errors + assert all(p.rows_complete for p in clean.partitions) and not clean.file_errors (tmp_path / "extra.csv").write_text("a,b\n1,2\n") with_csv = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) - assert all(p.stats_complete for p in with_csv.partitions) # the parquet rows are still complete + assert all(p.rows_complete for p in with_csv.partitions) # the parquet rows are still complete assert with_csv.file_errors # but there is data here that went unprofiled assert with_csv.sampling.files_read == 1 and with_csv.sampling.files_present == 2 @@ -842,7 +866,6 @@ def test_row_budget_is_divided_across_a_partitions_files(tmp_path): result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME, row_budget=400) assert result.sampling.rows_scanned == 400 # 400 / 4 files = 100 rows each - assert result.sampling.row_budget == 400 def test_rows_read_do_not_grow_when_a_dataset_is_resharded(tmp_path_factory): diff --git a/plugins/nemo-datasets/tests/test_profile_task.py b/plugins/nemo-datasets/tests/test_profile_task.py index 71d442e808..43fd69078e 100644 --- a/plugins/nemo-datasets/tests/test_profile_task.py +++ b/plugins/nemo-datasets/tests/test_profile_task.py @@ -10,7 +10,6 @@ import nemo_datasets_plugin.tasks.profile.run as run_mod import pyarrow as pa import pyarrow.parquet as pq -from nemo_datasets_plugin.profiler.pipeline import DEFAULT_ROW_BUDGET from nemo_platform import NeMoPlatform from nemo_platform_plugin.job_results import ResultRef from nemo_platform_plugin.jobs.constants import ( @@ -85,12 +84,13 @@ def test_task_passes_column_role_hints_to_the_profiler(tmp_path, monkeypatch): assert classification["dataset_type"] == "prompt_completion" -def test_task_defaults_to_the_profilers_row_budget(tmp_path, monkeypatch): +def test_task_reads_everything_by_default(tmp_path, monkeypatch): data = _dataset(tmp_path / "data") published = _install(monkeypatch, tmp_path, {"path": str(data)}) assert run_mod.run(_SDK) == 0 - assert published["profile"]["sampling"]["row_budget"] == DEFAULT_ROW_BUDGET + sampling = published["profile"]["sampling"] + assert sampling["rows_scanned"] == sampling["rows_present"] # nothing left unread def test_task_honours_an_explicit_row_budget(tmp_path, monkeypatch): @@ -98,7 +98,7 @@ def test_task_honours_an_explicit_row_budget(tmp_path, monkeypatch): published = _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": 5}) assert run_mod.run(_SDK) == 0 - assert published["profile"]["sampling"]["row_budget"] == 5 + assert published["profile"]["sampling"]["rows_scanned"] == 5 assert published["profile"]["sampling"]["rows_scanned"] == 5 @@ -107,8 +107,9 @@ def test_row_budget_zero_asks_for_every_row(tmp_path, monkeypatch): published = _install(monkeypatch, tmp_path, {"path": str(data), "row_budget": 0}) assert run_mod.run(_SDK) == 0 - assert published["profile"]["sampling"]["row_budget"] is None - assert published["profile"]["partitions"][0]["stats_complete"] is True + sampling = published["profile"]["sampling"] + assert sampling["rows_scanned"] == sampling["rows_present"] # 0 means "all of them" + assert published["profile"]["partitions"][0]["rows_complete"] is True def test_task_fails_when_the_step_config_says_nothing_to_profile(tmp_path, monkeypatch): From c19582274ef05c31ff8af476f3ef3c2a62f46c45 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Sat, 8 Aug 2026 00:11:54 -0400 Subject: [PATCH 44/44] fix(datasets): sample quality in blocks, and count a partial read for what it read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five findings from a review pass over the branch. Two were defects the tests did not cover, and one of them contradicts a claim I had written into a comment on the strength of too little measurement. **The quality sample aliased against periodic data.** It was taken at an even step, and data is periodic more often than it looks -- a set that round-robins over ten sources, or carries k responses per prompt, is periodic by construction. When the step shares a factor with that period it samples one phase and only that phase. At the shipped constants: 500,000 rows with every tenth corrupt gives a step of ten, and a reported repetition score of 1.000 against a truth of 0.100. Not noise. The wrong answer, tenfold. I had measured this before and cleared it. One dataset, HelpSteer2, period two, ≤8% drift -- and I wrote "does not buy a block-sampling scheme to avoid" into the comment. One sample whose period happened not to align with the step, generalised. Block sampling is exactly the fix I dismissed. The sample is now contiguous blocks of 512 rows, spaced evenly. A block longer than the period sees every phase of it, whatever the period is, and costs the same. Measured across periods 2, 5, 10, 25 and 100 at 500,000 rows: every one within 2%. The unknown-row-count path doubles the cycle rather than the step, so blocks stay whole as it thins, and the Horvitz-Thompson weight comes along unchanged. **A file that failed partway was counted as unread.** `rows_scanned` and `files_read` were incremented after the batch loop, which was right when a read was all-or-nothing and wrong once it streamed: a fold cannot give rows back, so batches already folded are in the statistics whatever happens next. A failure on the third batch reported `rows_scanned: 0` and `files_read: 0` beside stats built from 2,048 rows. Both now count what was actually consumed. **Three smaller things.** The stored contract pointed twice at `SplitProfile.files`, a field deleted several phases ago -- that is the docstring a consumer reads to understand the type. The relational prefix probe ran outside the per-column guard, so a failure in it would have surfaced as a `FileError`, collapsing the one distinction the two failure domains exist to keep. And `_PartitionFolds._declared` was assigned but never read, with `_per_file_cap` computed twice per partition. Both defects have regression tests written to fail against the previous code. The demo profile is byte-identical: these columns sit under the sample bound, so both schemes measure every row, and the fix only moves what the old one got wrong. Signed-off-by: Albert Cui --- .../files/dataset_profile.py | 19 +++++-- .../nemo_datasets_plugin/profiler/pipeline.py | 40 +++++++++---- .../nemo_datasets_plugin/profiler/stats.py | 57 ++++++++++++------- plugins/nemo-datasets/tests/test_pipeline.py | 27 +++++++++ plugins/nemo-datasets/tests/test_stats.py | 31 ++++++++-- 5 files changed, 134 insertions(+), 40 deletions(-) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py index 581ef74c48..5a2e247c7b 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/files/dataset_profile.py @@ -206,9 +206,12 @@ class TextQuality(BaseModel): tens of thousands of rows pins down far past the precision anyone reads them to. Bounding them is what makes reading every row of a dataset affordable. - The sample is evenly strided: deterministic, so two runs over the same bytes agree, and spread - across the column rather than taken from its head, so a sorted shard does not decide the answer. - A column smaller than the bound is measured in full. + The sample is contiguous blocks, spaced evenly across the column: deterministic, so two runs over + the same bytes agree, and spread rather than taken from the head, so a sorted shard does not + decide the answer. Blocks rather than every n-th row because an even step aliases against + periodic data — a set that round-robins over sources, or carries k responses per prompt, is + periodic by construction, and a step sharing a factor with that period samples one phase and + only that phase. A column smaller than the bound is measured in full. """ whitespace_ratio: float = Field(ge=0.0, le=1.0, description="Padding / bad scraping.") @@ -420,7 +423,8 @@ class SplitProfile(BaseModel): num_examples: int | None = Field( default=None, description=( - "Rows in this split, counting every file in `files` whether or not its rows were read. Always " + "Rows in this split, counting every one of its files whether or not that file's rows were " + "read. Always " "exact — summed from parquet footers or from files read to their end — and None the moment any " "one file's count is unknown. Never an estimate, so it carries no accuracy caveat: a capped run " "still reports the true total whenever the footers knew it." @@ -531,7 +535,12 @@ class SamplingInfo(BaseModel): ), ) files_read: int = Field( - description="Files actually opened and read from (a count; the files themselves are `SplitProfile.files`)." + description=( + "Files actually opened and read from. A count, not a list -- the paths of healthy " + "shards are the one part of a profile that grows without bound and informs no " + "decision; `SplitProfile.num_files` counts them per split, and only the ones that " + "went wrong are named, on `DatasetProfile.file_errors`." + ) ) files_present: int = Field( description=( diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py index 4199c2867f..1a37e8e319 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/pipeline.py @@ -18,7 +18,8 @@ What a declared schema buys is not the fold but its sharpness. Parquet footers are read first, so the columns are known before a row is parsed and each accumulator is chosen up front, and the exact -row count is known too, which is what lets a quality stride be placed over a column not yet seen. +row count is known too, which is what lets the quality sample be spread across a column not yet +seen. Without one — line-delimited data — both wait for the data. Columns are created on first sight and back-filled with the rows they were absent for, and each carries every shape at once until the last @@ -283,7 +284,7 @@ def _peek_files(source: FileSource, entries: list[FileEntry]) -> dict[str, FileP def _expected_rows(previews: dict[str, FilePreview], row_cap: int | None) -> int | None: """How many rows the fold is about to see, if every file said. - Capped per file the same way the read will be, so a budgeted run strides its quality sample over + Capped per file the same way the read will be, so a budgeted run spreads its quality sample over what it will actually scan rather than over what the dataset holds. """ total = 0 @@ -305,15 +306,25 @@ def __init__(self, features: list[FeatureSchema] | None, expected_rows: int | No # Declared: the columns are known, so the accumulators are chosen now. Inferred: they are # discovered as they appear and typed once every row has gone by. self.features = features or [] - self._declared = features is not None self._columns: ColumnFold | InferredColumnFold = ( ColumnFold(features, expected_rows) if features is not None else InferredColumnFold(expected_rows) ) self._prefix = PrefixPairFold() + self._prefix_error: Evidence | None = None def update(self, rows: list[dict]) -> None: self._columns.update(rows) - self._prefix.update(rows) + # Guarded like the columns are. Unguarded, the only thing that would catch this is the + # per-file handler, which would report odd *data* as a bad *file* -- collapsing the one + # distinction the two failure domains exist to keep. + if self._prefix_error is None: + try: + self._prefix.update(rows) + except Exception as exc: + self._prefix_error = Evidence( + kind="error", + detail=f"the chosen/rejected prefix probe could not run: {type(exc).__name__}: {exc}", + ) def measure( self, column_roles: dict[str, str] @@ -339,6 +350,8 @@ def measure( quote_enumerations(self.features, measured.stats, measured.vocabularies) classification.evidence.extend(_capped_columns_evidence(self.features)) classification.evidence.extend(measured.errors) + if self._prefix_error is not None: + classification.evidence.append(self._prefix_error) return self.features, measured.stats, classification except Exception as exc: detail = f"could not measure this partition: {type(exc).__name__}: {exc}" @@ -380,15 +393,15 @@ def _profile_partition( # Declared or not, the partition folds. With a schema the accumulators are chosen up front and # the exact row count places the quality stride; without one both wait for the data, which costs # a deferred dtype per column and nothing else. + row_cap = _per_file_cap(row_budget, len(entries)) folds = _PartitionFolds( derive_features([], declared) if declared is not None else None, - expected_rows=_expected_rows(previews, _per_file_cap(row_budget, len(entries))), + expected_rows=_expected_rows(previews, row_cap), ) rows_scanned = 0 files_read = 0 rows_present: int | None = 0 partition_scanned = True - row_cap = _per_file_cap(row_budget, len(entries)) file_errors: list[FileError] = [] file_formats: set[str] = set() split_profiles: list[SplitProfile] = [] @@ -401,18 +414,16 @@ def _profile_partition( error: str | None = None num_rows: int | None = None scanned_all = False + scanned = 0 try: # Inside the guard: resolving the reader can fail too, and a format with no reader # registered is a file the profiler could not use like any other. reader = get_reader(_format_of(entry.path)) preview = previews[entry.path] - scanned = 0 read_errors: list[str] = [] for batch in reader.batches(source, entry, row_cap=row_cap, errors=read_errors): folds.update(batch) scanned += len(batch) - files_read += 1 - rows_scanned += scanned # A file the reader only partly understood is named, the same as one it could not # open at all. Folding it silently would make a corrupt shard look complete. error = "; ".join(read_errors) or None @@ -427,11 +438,18 @@ def _profile_partition( scanned_all = num_rows is not None and scanned >= num_rows and error is None except Exception as exc: # Failure isolation: an unreadable file (or missing reader) keeps its identity, - # skips its rows, and does not abort the profile. The reason is recorded rather than - # swallowed, so a consumer can tell corrupt input from a profiler bug. + # skips the rest of its rows, and does not abort the profile. The reason is recorded + # rather than swallowed, so a consumer can tell corrupt input from a profiler bug. error = f"{type(exc).__name__}: {exc}" num_rows = None scanned_all = False + # Counted for what was actually consumed, outside the guard, because a read is no longer + # all-or-nothing: a fold cannot give rows back, so a file that failed on its fifth batch + # still contributed four and the envelope has to say so. Accounting for it as unread + # would leave `rows_scanned` describing fewer rows than the stats were built from. + rows_scanned += scanned + if scanned or error is None: + files_read += 1 if error is not None: file_errors.append(FileError(path=entry.path, error=error)) rows_present = _add_known(rows_present, num_rows) diff --git a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py index a999ad9d59..071765282a 100644 --- a/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py +++ b/plugins/nemo-datasets/src/nemo_datasets_plugin/profiler/stats.py @@ -21,10 +21,10 @@ partition before it can measure it. The base class is the entire measurement for a dtype with no statistics of its own, because the probes run over every column whatever its type. -Every measurement is now O(1) in rows except one: a string column retains its strings, because the -quality stride needs the column's row count to place its sample and that is not known until the last -batch. Summing the parquet footers before folding is what removes it — the row count is in them, and -reading it costs no rows. A messages column has no such term and is already bounded. +Every measurement is O(1) in rows. Nothing is retained: the one thing that used to be — a string +column's values, held so a quality sample could be placed across them — went when the sample learned +to place itself as it goes, in contiguous blocks whose size does not depend on how long the column +turns out to be. """ from __future__ import annotations @@ -357,24 +357,25 @@ def __init__(self, expected_rows: int | None = None) -> None: self._quality = _TextQualityCounters() self._seen = 0 self._sampled = 0 - # With the row count known -- parquet footers give it before a row is read -- the stride is - # fixed now and the sample is spread evenly over the whole column. Without it there is no - # length to stride over yet, so the stride starts at one and doubles each time the sample + # With the row count known -- parquet footers give it before a row is read -- the cycle is + # fixed now and the blocks spread evenly over the whole column. Without it there is no length + # to spread over yet, so the cycle starts at one block and doubles each time the sample # fills: every row eligible at first, thinning as the column turns out to be long. Each - # sampled row then stands for `stride` rows, which is what keeps the estimate unbiased - # rather than weighted toward the head where sampling was densest. - self._stride = _quality_stride(expected_rows) if expected_rows is not None else 1 + # sampled row then stands for `cycle / block` rows, which is what keeps the estimate + # unbiased rather than weighted toward the head where sampling was densest. + self._block = _QUALITY_SAMPLE_BLOCK + self._cycle = _quality_cycle(expected_rows) if expected_rows is not None else self._block self._adaptive = expected_rows is None def _observe(self, present: list[Any]) -> None: for value in present: if isinstance(value, str): self._lengths.add(len(value)) - if self._seen % self._stride == 0: - self._quality.add(value, self._stride) + if self._seen % self._cycle < self._block: + self._quality.add(value, self._cycle / self._block) self._sampled += 1 if self._adaptive and self._sampled >= _QUALITY_SAMPLE_ROWS: - self._stride *= 2 + self._cycle *= 2 self._sampled = 0 self._seen += 1 self._vocabulary.update(present) @@ -504,7 +505,8 @@ def _accumulator_for(feature: FeatureSchema, expected_rows: int | None = None) - """The accumulator that knows how to measure this column, dispatched once on its dtype. ``expected_rows`` is the partition's row count when it is known before reading -- only a string - column uses it, to place its quality stride without retaining the column. + column uses it, to space its quality blocks across the whole of itself rather than thinning + as it goes. """ if feature.dtype == "string": return StringAccumulator(expected_rows) @@ -759,6 +761,19 @@ def _at(self, percentile: int) -> int: # the precision anyone reads them to. Bounding them is what makes reading every row affordable. _QUALITY_SAMPLE_ROWS = 50_000 +# ...and the sample is taken in contiguous blocks of this many rows, not at an even step. +# +# A step aliases. Data is periodic more often than it looks -- a set that round-robins over ten +# sources, or carries k responses per prompt, is periodic by construction -- and a step that shares a +# factor with the period samples one phase and only that phase. Measured before this was blocks: +# 500,000 rows with every tenth corrupt gives a step of ten, which reported a repetition score of +# 1.000 against a truth of 0.100. Not noise; the wrong answer. +# +# A block longer than the period sees every phase of it, whatever the period is, and costs exactly +# the same. 512 covers anything plausible -- k-per-prompt is single digits, round-robin over sources +# is tens to low hundreds. +_QUALITY_SAMPLE_BLOCK = 512 + def _whitespace_count(text: str) -> int: """Whitespace characters, matching ``\\s`` exactly. @@ -787,7 +802,7 @@ def _non_ascii_count(text: str) -> int: class _TextQualityCounters: - """The three corruption ratios as running sums, so a strided sample needs no storage. + """The three corruption ratios as running sums, so a sampled subset needs no storage. Every denominator is the sample's own, never the column's: each ratio is an estimate over the rows actually scanned, which is what keeps it unbiased rather than diluted. @@ -800,7 +815,7 @@ def __init__(self) -> None: self._repetition = 0.0 self._rows = 0 - def add(self, text: str, weight: int = 1) -> None: + def add(self, text: str, weight: float = 1.0) -> None: """Fold one sampled row in, standing for ``weight`` rows of the column. The weight is what keeps a *varying* sample rate honest. Sampling one row in four and @@ -821,9 +836,13 @@ def finalize(self) -> TextQuality: ) -def _quality_stride(rows: int) -> int: - """How many rows to step between quality samples, given the column's length.""" - return max(1, rows // _QUALITY_SAMPLE_ROWS) +def _quality_cycle(rows: int) -> int: + """How many rows one sampled block stands in for, when the column's length is known. + + Every ``cycle`` rows, the first ``_QUALITY_SAMPLE_BLOCK`` of them are measured. Sized so the + blocks add up to the sample budget and spread across the whole column. + """ + return max(_QUALITY_SAMPLE_BLOCK, rows * _QUALITY_SAMPLE_BLOCK // _QUALITY_SAMPLE_ROWS) def _count_matches(pattern: re.Pattern[str], text: str) -> int: diff --git a/plugins/nemo-datasets/tests/test_pipeline.py b/plugins/nemo-datasets/tests/test_pipeline.py index 330928eb5c..5718dc6c8c 100644 --- a/plugins/nemo-datasets/tests/test_pipeline.py +++ b/plugins/nemo-datasets/tests/test_pipeline.py @@ -682,6 +682,33 @@ def poison_one_partition(features, stats, **kwargs): assert partitions["good"].stats # a neighbour's bad data costs this partition nothing +def test_a_file_that_fails_partway_still_counts_what_it_contributed(tmp_path, monkeypatch): + # A read used to be all-or-nothing, so a failure meant no rows at all and the envelope could be + # written after it. A fold cannot give rows back: batches already folded are in the statistics + # whatever happens next, and counting the file as unread left `rows_scanned` describing fewer + # rows than the stats were built from. + from nemo_datasets_plugin.profiler import pipeline as pipeline_module + + _write_parquet(tmp_path / "train.parquet", [{"a": i} for i in range(4000)]) + real_update = pipeline_module._PartitionFolds.update + calls = {"n": 0} + + def fail_on_the_third_batch(self, rows): + calls["n"] += 1 + if calls["n"] == 3: + raise RuntimeError("boom mid-file") + return real_update(self, rows) + + monkeypatch.setattr(pipeline_module._PartitionFolds, "update", fail_on_the_third_batch) + result = profile(LocalFileSource(tmp_path), created_at=FIXED_TIME) + + assert result.sampling.rows_scanned == 2048 # two batches of 1024 were folded before it failed + assert result.sampling.files_read == 1 # the file *was* read from, just not to its end + assert result.partitions[0].stats["a"].numeric is not None # and those rows shaped the stats + assert [e.path for e in result.file_errors] == ["train.parquet"] + assert result.partitions[0].rows_complete is False + + def test_reading_everything_is_the_default(tmp_path): # The point of the whole exercise. The budget existed to keep a materialised partition off the # heap; nothing is materialised, so the default should not answer the question worse than it can diff --git a/plugins/nemo-datasets/tests/test_stats.py b/plugins/nemo-datasets/tests/test_stats.py index 70a8fe105c..fe83c8a0bf 100644 --- a/plugins/nemo-datasets/tests/test_stats.py +++ b/plugins/nemo-datasets/tests/test_stats.py @@ -257,6 +257,26 @@ def test_quality_fast_paths_are_the_same_measurement_as_the_regexes(text): assert _non_ascii_count(text) == sum(1 for _ in _NON_ASCII_RUN.finditer(text)) +def test_the_quality_sample_does_not_alias_against_periodic_data(monkeypatch): + # A set that round-robins over sources, or carries k responses per prompt, is periodic by + # construction. An evenly-spaced step whose spacing shares a factor with that period samples one + # phase and only that phase: 500,000 rows with every tenth corrupt gave a step of ten and a + # repetition score of 1.000 against a truth of 0.100. A contiguous block longer than the period + # sees every phase of it, whatever the period turns out to be. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 1_000) + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_BLOCK", 64) + period, n = 10, 100_000 + values = ["aaaaaaaaaaaa" if i % period == 0 else "the quick brown fox" for i in range(n)] + + known = stats_module.StringAccumulator(n) + known.update(values) + assert known.finalize()[0].quality.repetition_score == pytest.approx(1 / period, abs=0.02) + + unknown = stats_module.StringAccumulator(None) + unknown.update(values) + assert unknown.finalize()[0].quality.repetition_score == pytest.approx(1 / period, abs=0.02) + + def test_a_known_row_count_strides_evenly_and_deterministically(monkeypatch): # With the row count known up front the stride is fixed, so the sample is spread evenly over the # whole column -- and two runs over the same bytes agree, which is why no RNG is involved. @@ -274,17 +294,18 @@ def quality(expected_rows): def test_an_unknown_row_count_thins_as_it_goes_and_stays_unbiased(monkeypatch): - # No footer, so no length to stride over: the stride starts at one and doubles as the sample - # fills. Sampling is then densest at the head, which would skew the answer -- weighting each - # sampled row by the stride it stood for is what corrects it. - monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 10) + # No footer, so no length to spread blocks over: the cycle starts at one block and doubles as the + # sample fills. Sampling is then densest at the head, which would skew the answer -- weighting + # each sampled row by the rows its block stood for is what corrects it. + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_ROWS", 40) + monkeypatch.setattr(stats_module, "_QUALITY_SAMPLE_BLOCK", 8) values = ["clean text"] * 500 + ["aaaaaaaaaaaa"] * 500 acc = stats_module.StringAccumulator(None) acc.update(values) score = acc.finalize()[0].quality.repetition_score - assert acc._stride > 1 # it did thin + assert acc._cycle > acc._block # it did thin assert 0.35 <= score <= 0.65 # ...and still found roughly half the column corrupt