Skip to content

feat: Pydantic RNTuple extension and Iterative/FuturesBackends - #1579

Draft
NJManganelli wants to merge 9 commits into
scikit-hep:masterfrom
NJManganelli:pydantic_claude_rev_and_rntuple
Draft

feat: Pydantic RNTuple extension and Iterative/FuturesBackends#1579
NJManganelli wants to merge 9 commits into
scikit-hep:masterfrom
NJManganelli:pydantic_claude_rev_and_rntuple

Conversation

@NJManganelli

@NJManganelli NJManganelli commented Jul 4, 2026

Copy link
Copy Markdown
Collaborator

Adds RNTuple capability to preprocessing, plus alternative backends (besides dask) to permit different processing modes. The backend design is intended to potentially interface with the coffea.compute refactor in the future.

This has been largely implemented by Claude Fable 5 + Opus 4.8, with hands-on guidance from me

This PR is rebased to #1528

Todos:

  • cleanup filespec.ipynb to better demonstrate the method-chaining for slicing etc
  • decide on whether this backends approach makes sense right now
  • full human review of code

@NJManganelli NJManganelli changed the title Pydantic RNTuple extension and Iterative/FuturesBackends feat: Pydantic RNTuple extension and Iterative/FuturesBackends Jul 4, 2026
@NJManganelli
NJManganelli force-pushed the pydantic_claude_rev_and_rntuple branch 2 times, most recently from c27d8d9 to 4e883a8 Compare July 20, 2026 03:00
Nick Manganelli and others added 9 commits July 20, 2026 00:38
…etection

- get_steps now auto-detects TTree vs RNTuple per file: RNTuple num_entries,
  cluster boundaries (from cluster_summaries), and form (via uproot.dask with the
  {path: object_path} spec, since the opened RNTuple object is not accepted) are
  handled alongside the existing TTree path. preprocess()/preprocess_root handle
  RNTuples transparently.
- Add preprocess_rntuple(), mirroring preprocess_parquet/preprocess_root, which
  requires every object to be an RNTuple (require_rntuple=True) and raises
  otherwise; exported from coffea.dataset_tools.
- Extract shared helpers _even_steps, _aligned_steps, _rntuple_cluster_boundaries,
  and _union_form_jsonstr; use them in get_steps and get_parquet_form_uuid_steps
  and de-duplicate the union-form logic across the legacy and pydantic paths.
- Short-circuit an empty per-format DataGroupSpec in _preprocess_pydantic so a
  single-format fileset no longer issues a wasted (empty) dask.compute.
- Remove the spurious dask function-cache pop in the parquet path (no graph is
  built there) and guard the TTree/RNTuple pop against an empty cache; drop the
  always-true 'num_entries >= 0' guard; correct file_exceptions/parquet docstrings.
- Add end-to-end test_preprocess_rntuple and a filter_files->PreprocessedFiles test.

Assisted-by: Claude Opus 4.8
…ree forms

Factor the per-dataset preprocessing map-reduce out of _preprocess_pydantic
behind a small backend interface (new module dataset_tools/backends.py) shaped
to align with the coffea.compute refactor (PR scikit-hep#1470): PreprocessJob (computable)
-> Backend.submit() -> future-like Task.result(), with an ordered concat reduce
(preprocessing zips results against original file order).

Backends:
- DaskBackend: the historical from_awkward/map_partitions/AwkwardTreeReductionLayer
  graph; the only backend that imports dask for orchestration.
- IterativeBackend: immediate (synchronous, single process), mirroring
  IterativeExecutor. Dask-free.
- FuturesBackend: concurrent.futures pool, threads by default with a use_processes
  opt-in. Dask-free.

preprocess()/preprocess_root/_rntuple/_parquet gain backend= (default "dask", so
existing behavior is unchanged). On a dask import failure the default path prints
a coffea_console hint pointing at the dask-free backends. The Runner is not
rewired; that is left to the compute refactor.

TTree form extraction is now dask-free: it uses uproot's own non-dask form builder
(uproot._dask._get_ttree_form, guarded with a fallback), byte-identical to
uproot.dask(tree).layout.form, which also drops the function_cache workaround for
TTrees. RNTuple form extraction still uses uproot.dask (documented).

Also fixes found while self-reviewing this change:
- get_parquet_form_uuid_steps: add the num_entries==0 guard ROOT already had, so a
  0-row parquet yields steps=[[0,0]] instead of ZeroDivisionError.
- Scope the dask fallback hint to submit() import failure so a worker's own
  ModuleNotFoundError during compute is no longer misreported as "dask missing".
- FuturesBackend: shut down an owned pool via __del__ if the task is only wait()-ed
  or dropped without result().
- resolve_backend: warn when scheduler is passed with a pre-built backend instance
  (previously dropped silently for a DaskBackend instance).
- DaskBackend: clamp files_per_batch with max(1, ...) for parity with the other
  backends.

Tests: tests/test_dataset_tools_backends.py; dask/iterative/futures verified to
produce equal DataGroupSpec, and the non-dask TTree form verified byte-identical
to the dask form.

Assisted-by: Claude Opus 4.8
An explicit step_size=0 (or negative) previously surfaced as a bare
ZeroDivisionError from _even_steps deep inside a worker/dask graph. Add a
_validate_step_size guard, called at the start of _preprocess_pydantic (covers
preprocess/preprocess_root/_rntuple/_parquet) and preprocess_legacy, raising a
clear ValueError before any file I/O. None and >= 1 are unaffected.

Assisted-by: Claude Opus 4.8
Extend the dask-free form path from TTree to RNTuple by routing both through
uproot's own form builder (uproot._dask._get_ttree_form, which already handles
RNTuple HasFields sub-branches) and nulling the resulting form keys. The RNTuple
delay-open dask path builds base_form the same way; the only delta was that
uproot.dask exposes the keyless meta form while _get_ttree_form carries RNTuple
"column-N" keys, so _null_form_keys closes the gap (a no-op for already-keyless
TTree forms).

_ttree_form_json is replaced by the format-aware _awkward_form_json (filter_field
+ full_paths for RNTuple, filter_branch for TTree). uproot.dask is now only a
fallback for the unlikely case that the private uproot helper is unavailable. Net
result: iterative/futures backends are fully dask-free for parquet and both ROOT
flavors.

Verified byte-identical to uproot.dask(...).layout.form for TTree and RNTuple
samples (parametrized test), and dask/iterative/futures produce equal
DataGroupSpec end-to-end for an RNTuple dataset.

Assisted-by: Claude Opus 4.8
… rename backends module

Pre-pay the coffea.compute (PR scikit-hep#1470) migration on two axes:

(a) Order-relaxation. _preprocess_pydantic previously zipped the concatenated
    per-file results positionally against the original input to recover skipped/
    bad files (emitted as None by the worker), which made the reduce order
    load-bearing. Assemble both `available` and `updated` filesets by filename
    instead, in the original input order, so a skipped file is identified by its
    absence from the results rather than by position. Correctness no longer
    depends on the order in which a backend reduces/concatenates, which is what
    lets preprocessing mesh with an order-agnostic compute reduce. The shared
    get_steps worker (also used by the legacy path) is untouched.

(b) Task.partial_result(). Add partial_result() to the PreprocessTask protocol
    and all three tasks, mirroring compute.Task: the eager task returns its full
    result, the futures task gathers only completed batches without blocking (and
    without shutting the pool), and the dask task falls back to the full result
    (a fused graph has no cheap partial).

Also rename backends.py -> preprocess_backends.py (clearer alongside a future
coffea.compute.backends) and soften the now-obsolete "ordering is load-bearing"
docs.

Tests: filename-based assembly of a skipped bad file; partial_result == result
when complete. Full preprocess suite + backend suite green; pre-commit clean.

Assisted-by: Claude Opus 4.8
…ocessing

filespec.ipynb: extend the "Integration with Preprocessing" section with three
new subsections (executed against the on-master samples, real outputs attached):
- 5.2 Switchable execution backends (backend="dask"/"iterative"/"futures",
  FuturesBackend instances), showing the dask-free backends agree with dask.
- 5.3 Preprocessing Parquet datasets (preprocess_parquet, format auto-detection,
  no object_path, use_row_groups).
- 5.4 Preprocessing RNTuple datasets (auto-detection; preprocess_rntuple's
  RNTuple-only contract).
Imports and the overview are updated accordingly; later cells' execution counts
are renumbered to stay monotonic. The remaining cells' committed outputs are
untouched (the apply_to_fileset/dask-histogram compute path can't be re-run in
the current env due to a known awkward/dask-histogram version-compat baseline,
unrelated to preprocessing).

processing.ipynb: add a short note to the Preprocessing section pointing at the
new backend= argument and Parquet/RNTuple support, linking to filespec.ipynb.

(docs/source/notebooks/*.ipynb are symlinks into binder/, so they update too.)

Assisted-by: Claude Opus 4.8
Re-execute the entire notebook end-to-end (nbconvert --execute) now that the
environment is unblocked: apply_to_fileset -> dask.compute runs cleanly on the
advanced scikit-hep stack (uproot 5.7.4, awkward 2.10.0, dask-awkward 2026.2.1,
dask-histogram 2026.2.0), which previously crashed with the awkward/dask-awkward
scalar-meta incompatibility. Every cell (including the Integration-with-
apply_to_fileset section and the 5.2/5.3/5.4 backends/parquet/rntuple demos) now
carries real, consistently-numbered outputs from a single fresh run -- replacing
the earlier partial-output render. 0 error outputs.

Assisted-by: Claude Opus 4.8
…bsent

test_awkward_form_json_matches_uproot_dask compares the dask-free form
builder against uproot.dask output, so it requires dask and dask-awkward;
guard it with pytest.importorskip like the other dask-dependent tests in
this module so the no-dask CI job skips it.

Assisted-by: Claude Fable 5
- require_rntuple raises inside the per-file error handling so
  skip_bad_files/file_exceptions can skip TTree files, and rejects
  parquet-format datasets with a clear error
- preprocess_root gains require_rntuple=; preprocess_rntuple becomes a
  thin alias instead of a duplicated 13-parameter signature
- get_steps/get_parquet_form_uuid_steps validate step_size directly
  (negative values previously produced a silent single step per file)
- drop the unreachable duplicate zero-row guard in
  get_parquet_form_uuid_steps
- forward uproot_options to the RNTuple uproot.dask form fallback
- futures backend: fail fast on batch failure (cancel pending batches),
  and default pool sizing to the executor default instead of 1 worker
- resolve_backend injects scheduler into a DaskBackend instance that has
  none; preprocess resolves the backend once so mixed filesets do not
  double-warn
- serialize the null-key form dict directly (skips a Form round-trip)
  and reuse a single to_list in the by-filename assembly
- comment-convention cleanups (no PR/review references, present-tense
  descriptions)

Assisted-by: Claude Fable 5
@NJManganelli

NJManganelli commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator Author

🤖 AI text below 🤖

Follow-up: exercising the RNTuple preprocessing + switchable backends with the bundled test files.

This PR adds RNTuple support (auto-detected, plus a strict preprocess_rntuple) and dask-free iterative/futures backends. Here's a minimal way for a maintainer to exercise both against the repo's own samples, going through the top-level DataGroupSpec API (no hand-built filespecs):

from coffea.dataset_tools import DataGroupSpec, preprocess, preprocess_rntuple

# One group mixing a TTree dataset and an RNTuple dataset; preprocess() auto-detects
# TTree vs RNTuple per file. backend="iterative" runs dask-free (single process).
fileset = DataGroupSpec({
    "ttree_dy":   {"files": {"tests/samples/nano_dy.root": "Events"}},
    "rntuple_dy": {"files": {"tests/samples/nano_dy_rntuple.root": "Events"}},
})
available, _ = preprocess(fileset, step_size=15, backend="iterative")
for name, ds in available.items():
    fs = next(iter(ds.files.values()))
    print(name, fs.steps, fs.num_entries)
# ttree_dy   [[0, 14], [14, 28], [28, 40]] 40
# rntuple_dy [[0, 14], [14, 28], [28, 40]] 40   <- identical steps/form via RNTuple path

# preprocess_rntuple asserts every object is an RNTuple (alias for
# preprocess_root(..., require_rntuple=True)); a TTree in the group would raise
# a clear ValueError (subject to skip_bad_files/file_exceptions). Here on a thread pool:
rnt_only = DataGroupSpec({"rn": {"files": {"tests/samples/nano_dimuon_rntuple.root": "Events"}}})
avail_rnt, _ = preprocess_rntuple(rnt_only, backend="futures")
print(list(avail_rnt["rn"].files))

The iterative and futures backends produce a DataGroupSpec byte-identical (decoded-form-wise) to the dask backend, including the saved union form, for both TTree and RNTuple — so this runs with or without the dask stack installed. tests/test_dataset_tools_backends.py pins that equivalence (and tests/test_dataset_tools.py::test_preprocess_rntuple pins the RNTuple form against uproot.dask).

Three follow-on features build on this branch as stacked draft PRs against master (each includes this branch's commits until it merges, then collapses to just its own feature): union forms via DatasetSpec addition #1599 (#1478), per-file metadata extraction #1600 (#1401), and a mutable-steps prototype #1601 (#1483).


Developed with Claude Fable 5.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant