feat: Pydantic RNTuple extension and Iterative/FuturesBackends - #1579
feat: Pydantic RNTuple extension and Iterative/FuturesBackends#1579NJManganelli wants to merge 9 commits into
Conversation
c27d8d9 to
4e883a8
Compare
…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
4e883a8 to
327ebca
Compare
|
🤖 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 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 Three follow-on features build on this branch as stacked draft PRs against Developed with Claude Fable 5. |
327ebca to
3e9788e
Compare
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: