feat: add hooks for working with graphed-org/graphed

[lgray]

Add hooks to uproot that allow it to be used with the delayed execution engine "graphed", it has stabilized enough to where it makes sense to support the code upstream and ease installation.

🤖 AI text below 🤖

Added 2026-09-19 — uproot.graphed takes the same form-mapping seam uproot.dask does. Three parameters, each mirroring its uproot.dask namesake so a schema library (coffea's NanoEventsFactory is the archetype) can drive either arm from one mapping:

• form_mapping= — the ImplementsFormMapping object uproot.dask takes. The recorded array has the MAPPED form, the mapping info's behavior is registered on the backend, and the source declares its read list in TBranch names, so projection still fetches only the branches the recorded graph touches.
• known_base_form= — use this form instead of opening a file to learn it; no file is opened at record time, and the TTree name comes from the object path in files.
• backend= — a graphed backend instance to record on, in place of the default graphed.awkward.AwkwardBackend (the route for a caller's own Array subclass); behavior= is refused beside it.

Along the way: every partition is opened through the source, so the caller's file options reach both read and write; an empty declared read list is honoured; an RNTuple's base form is built the way a TTree's is; and uproot.extras.graphed() registers uproot as graphed-internal, so a recorded op's provenance (and a run-time error) points at the analyst's line rather than an Uproot frame. That last call needs graphed>=0.0.3 (on PyPI since 2026-09-19), so the test group floor moves to 0.0.3 and extras.graphed() tells a user on an older graphed how to upgrade. Tests: tests/test_1720_graphed_form_mapping.py, …_mapped_projection.py, …_mapped_write.py, …_mapped_edge_cases.py over a small NanoAOD-shaped helper in tests/graphed/mini_nanoaod.py; src/uproot/_graphed.py stays at 100 % line+branch under the suite.

The graphed hooks ride uproot's own CI: graphed, graphed-executors and graphed-histogram (all on PyPI; graphed at 0.0.3, the other two at 0.0.2 or later) are entries in the test dependency group behind python_version >= "3.11" (graphed requires 3.11+), so the tests/test_1720_graphed_*.py suite runs on the 3.11+ jobs and is skipped on 3.10. The optional imports go through uproot.extras.graphed() / uproot.extras.graphed_executors(); the suite's non-test helpers live in the tests/graphed package and are imported as tests.graphed.<name>. The three packages are opted out of [tool.uv] exclude-newer the way awkward already is, because the 7-day window would otherwise hide fresh releases of the three; those keys can go once the window has rolled past.

The pre-commit fixes are committed by hand because pre-commit.ci cannot push to a fork owned by an organization. pyproject.toml keeps the layout pre-commit.ci's cached toml-sort 0.24 produces; a fresh local hook env pulls toml-sort 0.25.0 (2026-09-12), which reformats inline tables and would fail the same hook on pre-commit.ci. The "GPU test with CUDA 13" failure is the self-hosted runner's cuda-bindings pre-release version parse in tests-cuda/, which this PR does not touch and which main hits too.

[lgray]

@ariostas this one is ready to go now.

[ariostas]

Hi @lgray, I left a few comments. I'll take a closer look when I have some time.

[lgray]

@ariostas took care of your comments - but it looks like whatever serves the root files in the tests is 403'ing hard. I'll try kicking the tests again tomorrow.

[lgray]

@ariostas something's really wrong with the scikit-hep hosting all the CI jobs keep timing out.

[ariostas]

@ariostas something's really wrong with the scikit-hep hosting all the CI jobs keep timing out.

Hmm I'll have a look

[ariostas]

@lgray should work now. Thanks for letting me know! I restarted the CI

[lgray]

@ariostas cool - all green here now except CUDA-13 as expected. Lemme know if you have any comments :-)

[lgray]

@ariostas a gentle ping on this one when you have time.

[ariostas]

I did some skimming, but since it's a big PR, my review was mostly AI-driven. I left some comments below. This was with Fable. I'll see if a GPT model has any extra recommendations.

🤖 AI text below 🤖

Thanks for this — the projection story and the form-mapping seam are nicely done, and the suite is green and lint-clean locally. Since the goal is for this to eventually stand in for uproot.dask, I checked it against uproot.dask / dask_write behaviour and found a few correctness problems and parity gaps (inline). Three overall requests:

1. Mark it experimental. The top-level API depends on a 0.0.x package whose required release predates this PR update by hours. Please say so in the uproot.graphed / uproot.graphed_write docstrings (API may change without a deprecation cycle), and keep the public surface to those two names (+ maybe graphed_head) until it settles.
2. Collapse the two read paths into one by defaulting to TrivialFormMapping, as uproot.dask does. I tried it: 128/135 tests still pass, and the 7 that don't either assert unmapped-only reporting or depend on the bare-expression write fallback, which is a bug.
3. Trim the tests. On the ubuntu 3.11 CI job the test_1720_* block is 42 s of a 160 s pytest run (26 % of the time for 13 % of the tests); locally it is 23 s wall but 43 s CPU because of process pools, which is what a single-core runner pays. Almost half of that is four files that exercise only graphed_partitions / read_graphed_partition.

[ariostas]

🤖 AI text below 🤖

Bug — a bare expression writes the wrong data. graphed_write(ev.px1 * 2, dest) succeeds and writes unscaled px1: for a fieldless result this falls back to the raw input chunk. Through a form_mapping the same call crashes with AttributeError: 'PlaceholderArray' object has no attribute 'reshape', because the chunk holds placeholders for unread branches. Please raise a TypeError asking for a record (or write under an explicit default branch name) instead of substituting the inputs. test_write_of_a_projected_array_writes_only_its_branches currently pins the wrong behaviour (g.x + g.y → branches x, y).

[lgray]

🤖 AI text below 🤖

Fixed (W1). A fieldless result is refused with TypeError before any file is opened, telling the user to wrap it (gak.zip({"x2": g.x * 2})); nothing substitutes the inputs any more. Witness: test_derived_record_is_written_and_a_fieldless_expression_is_refused. The test that pinned the fallback is gone with the old suite.

[ariostas]

🤖 AI text below 🤖

Output type differs from dask_write. Dict assignment writes an RNTuple here, while dask_write calls mktree and writes a TTree — under a parameter called tree_name. For a drop-in replacement I'd default to mktree and make RNTuple an explicit opt-in.

[lgray]

🤖 AI text below 🤖

Fixed (W2). Each partition is written by _dask_write.ak_to_root, so the writer is mktree and the output is a TTree under tree_name, as dask_write does (test_writes_one_ttree_file_per_partition). I did not add an RNTuple opt-in: dask_write has none either, so there is no parity to keep and nobody has asked for it yet.

[ariostas]

🤖 AI text below 🤖

compression="ZLIB" (accepted by dask_write) and unknown names both raise a bare KeyError. ak_to_root in _dask_write.py already has this mapping with a proper error message — could the two share one helper?

[lgray]

🤖 AI text below 🤖

Fixed (W3). Compression goes through the same ak_to_root lookup: "ZLIB" is accepted and an unknown name gets its error message (test_compression_names_match_dask_write).

[ariostas]

🤖 AI text below 🤖

executor="threads" is a bare KeyError, and an executor instance fails with 'ThreadExecutor' object is not callable. A ValueError listing the accepted strings, and accepting an instance (use it as-is when it has .run), would be friendlier.

[lgray]

🤖 AI text below 🤖

Fixed (W4). executor= takes "process", "thread" or "sequential" (unknown string → ValueError listing them), an executor instance used as-is, or an executor class instantiated with max_workers (test_executor_is_a_name_an_instance_or_a_class).

[ariostas]

🤖 AI text below 🤖

Missing relative to dask_write: storage_options / remote destinations (os.makedirs makes this local-only, dask_write goes through url_to_fs), title, field_name, counter_name, initial_basket_capacity, resize_factor. Fine to defer, but then please list them as not-yet-supported in the docstring.

[lgray]

🤖 AI text below 🤖

Fixed (W5), not deferred. graphed_write now takes title, field_name, counter_name, initial_basket_capacity, resize_factor, compression_level and storage_options, all forwarded to ak_to_root; the destination is created and addressed through fsspec.core.url_to_fs, so remote destinations work and return full URLs while local ones return plain paths (test_tree_options_reach_mktree, test_remote_destinations_go_through_fsspec).

[ariostas]

🤖 AI text below 🤖

Putting these in the main test group runs the graphed suite on every 3.11+ job. Given the cost below, would a separate test-graphed group with one dedicated CI job (like test-xrootd / test-pyodide) work for you while this is experimental?

[lgray]

🤖 AI text below 🤖

Done (P1). The graphed test dependencies live in a test-graphed group (graphed[awkward,parquet]>=0.0.4, graphed-executors>=0.0.3, vector, all python_version >= '3.11'), and a single graphed-build job installs it. That job is red on this head only because the 0.0.4 floor is not on PyPI yet; graphed 0.0.4 (graphed-org/graphed#39) is being released now, after which the job gets re-run.

[ariostas]

🤖 AI text below 🤖

This is the most expensive file (9 s on CI): 20 parametrisations, the valid ones each spawning a 4-worker process pool, to cover the ~20-line range arithmetic in graphed_partitions. Asserting on the returned entry ranges directly (no executor) covers the same uproot code in milliseconds; one executor run is enough as an integration check. More generally, max_workers=4 process pools on a single-core runner mostly measure worker start-up — executor="thread" or a sequential runner would do for everything except one process-pool smoke test.

[lgray]

🤖 AI text below 🤖

Done (T1). The range arithmetic is asserted directly on the returned partitions, no executor (test_steps_per_file_splits_into_contiguous_ranges, test_step_size_caps_each_chunk). Everything else runs on the thread or sequential executor; one process-pool smoke test remains. The five files run in ~5 s wall locally.

[ariostas]

🤖 AI text below 🤖

In-process, this file, test_1720_graphed_report.py and test_1720_graphed_executor.py execute only graphed_partitions, read_graphed_partition and _partition_range from uproot; the assertions are about graphed's Bundle / DurablePlan / dead-letter / executor machinery. Those belong in graphed's own test suite — here they tie uproot CI to a 0.0.x package's behaviour. I'd drop the three files (and tests/graphed/report.py).

[lgray]

🤖 AI text below 🤖

Done (T2). test_1720_graphed_preserve.py, test_1720_graphed_report.py, test_1720_graphed_executor.py and tests/graphed/report.py are deleted; graphed's own machinery is tested in graphed.

[ariostas]

🤖 AI text below 🤖

This pins the fallback bug noted in _graphed_write.py: writing g.x + g.y is expected to produce branches x and y rather than the sum.

[lgray]

🤖 AI text below 🤖

Done (T3). This test is gone; the behaviour it pinned is now a TypeError (see the _graphed_write.py thread).

[ariostas]

🤖 AI text below 🤖

As you noted, this package shadows the real graphed when pytest runs from tests/. Renaming it (e.g. tests/graphed_helpers) removes the hazard for free. Also, 20 test files is a lot for one feature — after the cuts above, consolidating into a handful (read, form_mapping, write, projection) would make this much easier to maintain.

[lgray]

🤖 AI text below 🤖

Done (T4). The package is tests/graphed_helpers and the twenty files are consolidated into test_1720_graphed_read.py, _form_mapping.py, _projection.py, _write.py and _partitions.py (54 tests).

[ariostas]

🤖 AI text below 🤖

A GPT second pass found these additional, non-duplicate issues:

1. P1: nested RNTuples get the wrong schema (src/uproot/_graphed.py:316-341). HasFields.keys(full_paths=False) flattens rec.x to x, then _get_ttree_form advertises x as a top-level field even though RNTuple.arrays(["x"]) returns it under rec. In a local {rec: {x, y}, z} RNTuple, g.x records but fails at execution, and graphed_head(g.rec.x, 3) also fails. Please use to_akform(filter_name=common_keys) as _dask.py:1767-1769 does.

2. P1: repeated inputs overwrite output parts (src/uproot/writing/_graphed_write.py:181-199). file_bases is a dict keyed only by (uri, tree), so graphed([spec, spec]) with two steps reports part-00002.root, part-00003.root twice and writes only one copy of the rows. The occurrence needs to be part of the partition identity/index, or duplicate inputs need to be rejected.

3. P1: writes silently change axis-0 semantics with partition count (src/uproot/writing/_graphed_write.py:80-82). The complete expression is replayed independently on each chunk. For g.x[2:8] over range(10), one step writes [2,3,4,5,6,7], two steps write [2,3,4,7,8,9], and three write [2,5,8,9]; dask_write preserves dataset-global slicing. Either global operations need partition-aware handling or this path must reject/document non-row-local graphs.

4. P1: graphed_head applies n before the expression (src/uproot/_graphed.py:665-670). It reads the first n source rows and only then filters/slices. Thus graphed_head(g.x[10:20], 5) is empty instead of [10,11,12,13,14], and head(g[g.x % 2 == 0], 2) returns one row instead of two. A dask-like head must collect the first n output rows.

5. P1: mapped projection can under-supply compiled evaluation (src/uproot/_graphed.py:632-638). The comment above _evaluation_columns correctly notes that compiled IR replays dead intermediate nodes, but mapped sources switch back to output-buffer projection. With zip({"jet": g.Jet.pt, "muon": g.Muon.pt})[["jet"]], projection omits Muon_pt while IR still evaluates that node; both graphed_head and graphed_write fail. Mapped reads need the same syntactic-evaluation safeguard translated through the mapping.

6. P2: the default process writer cannot ship common behaviors (src/uproot/writing/_graphed_write.py:185-195). The partial serializes session.backend; vector.backends.awkward.behavior contains lambdas, so even the returned compute=False plan raises PicklingError under pickle.dumps. The generic Parquet writer already solves this with an importable behavior/backend reference; ROOT writing needs the same seam or should avoid serializing the backend instance.

7. P2: external evaluators are dropped by graphed_head and graphed_write (src/uproot/_graphed.py:663-670, src/uproot/writing/_graphed_write.py:82). A graphed.apply(...) expression materializes and runs through aggregate_plan, but these helpers call evaluate_ir without the session evaluators and fail with “External payload ... needs an evaluator.” aggregate_plan already wires these through _external_evaluators.

8. P2: already-open TTrees are not partition-compatible (src/uproot/_graphed.py:82-100). They are accepted by the documented dask-compatible files API and whole-dataset materialization works, but the Model_TTree is stored as Partition.uri; resources.open_once then raises because it is unhashable. Normalize it to a reusable file/tree reference or special-case open objects in open_tree.

The existing Graphed-focused suite still passes (135 passed); these are missing edge/parity cases rather than current test failures.

[lgray]

@ariostas thanks for spending the tokens - I will do so similarly once my limits reset :-)

[lgray]

🤖 AI text below 🤖

Pushed as four commits (dd14af0, 8565546, 5aac34e, bbb884b), one per class of finding. Three of the round-3 items are defects in graphed's own partition-wise drivers, so they are fixed there: graphed-org/graphed#38 (merged), and this PR's floor moves to graphed>=0.0.4, released as graphed-org/graphed#39. Until 0.0.4 is on PyPI the new graphed-build CI job cannot resolve the dependency; nothing else changes for the non-graphed jobs.

Every review item has its own reply: the round-2 inline comments are answered in their threads, and the round-2 overall requests and the round-3 items are answered in the quoted comments below. Three items are refuted or narrowed with evidence rather than fixed: round-3 item 5 (mapped projection — the example is not a valid awkward expression), and parts of G1 (the two executors were never dropped) and G6 (construction-time allow_missing already worked). Round-3 item 1's schema is fixed; the remaining nested-RNTuple read failure is uproot.dask's own and is described in that reply.

Also run locally: uproot's offline suite and the dask tests against the shared helpers (all green).

[lgray]

🤖 AI text below 🤖

Mark it experimental. … keep the public surface to those two names (+ maybe graphed_head) until it settles.

Done. Experimental notes on the module and on uproot.graphed / uproot.graphed_write; the exported names are uproot.graphed, uproot.graphed_head, uproot.graphed_write. The rest is importable from uproot._graphed only (details in the __init__.py thread).

[lgray]

🤖 AI text below 🤖

Collapse the two read paths into one by defaulting to TrivialFormMapping, as uproot.dask does.

Done, and further: reads share _dask.py's resolver, base-form construction and _read_tree with uproot.dask, so the two front ends cannot drift apart again (details in the _graphed.py L118 / L132 threads).

[lgray]

🤖 AI text below 🤖

Trim the tests. … Almost half of that is four files that exercise only graphed_partitions / read_graphed_partition.

Done. 20 files → 5 (read / form_mapping / projection / write / partitions), 54 tests, ~5 s wall locally; thread and sequential executors plus one process-pool smoke; graphed-machinery tests dropped; the graphed dependencies moved to a test-graphed group with one CI job.

[lgray]

🤖 AI text below 🤖

P1: nested RNTuples get the wrong schema … Please use to_akform(filter_name=common_keys) as _dask.py:1767-1769 does.

Fixed as suggested: the base form comes from to_akform(filter_name=common_keys) through the helper uproot.dask uses, so {rec: {x, y}, z} keeps its nesting and g.rec.x records (test_a_nested_rntuple_keeps_its_nesting).

One thing this does not fix, because it is not this PR's: reading a nested RNTuple field still fails, identically on uproot.dask — uproot.dask(f).rec.x.compute() raises ValueError: zip() argument 2 is longer than argument 1 from the RNTuple top-level-fields workaround in TrivialFormMappingInfo.load_buffers (the TODO there). Top-level RNTuple fields read fine on both paths.

[lgray]

🤖 AI text below 🤖

P1: repeated inputs overwrite output parts … The occurrence needs to be part of the partition identity/index, or duplicate inputs need to be rejected.

Rejected explicitly. graphed.write.file_bases (graphed-org/graphed#38) raises ValueError("duplicate input …: list each file (or file, tree) once"), and graphed_write surfaces it before any file is opened (test_refusals_happen_before_anything_is_written).

[lgray]

🤖 AI text below 🤖

P1: writes silently change axis-0 semantics with partition count … Either global operations need partition-aware handling or this path must reject/document non-row-local graphs.

Rejected explicitly, at the cause. A partition-wise driver evaluates the compiled IR once per chunk, so any reduction of the partitioned axis (x[2:8], x[0], gak.sum(x)) is a per-chunk partial. graphed.refuse_chunk_partials (graphed-org/graphed#38) refuses one that is a write output or that feeds another node, with GraphedError: 'slice' reduces the partitioned axis: a partitioned write has no combine step …; graphed_write calls it before opening anything (test_refusals_happen_before_anything_is_written). graphed's own to_parquet writers had the same defect and carry the same guard now. Dataset-global slicing is not emulated.

[lgray]

🤖 AI text below 🤖

P1: graphed_head applies n before the expression … A dask-like head must collect the first n output rows.

Fixed. graphed_head collects the first n rows of the output, growing a prefix read until the evaluated result has n rows or the first tree is exhausted: graphed_head(g.x[3:8], 3) == [3, 4, 5], graphed_head(g[g.x % 2 == 0], 2) returns two rows (test_head_returns_the_first_rows_of_the_output).

[lgray]

🤖 AI text below 🤖

P1: mapped projection can under-supply compiled evaluation … With zip({"jet": g.Jet.pt, "muon": g.Muon.pt})[["jet"]], projection omits Muon_pt while IR still evaluates that node; both graphed_head and graphed_write fail.

Not an under-supply. The earlier version of this reply said "not a valid awkward expression", which was imprecise — corrected here with measurements (same file, the MiniNanoMapping test helper, uproot.dask vs uproot.graphed).

The zip is pass-through for projection on both engines, and the read set is the accessed members plus the list structure the broadcast needs:

path	necessary columns	compute
uproot.dask, mapped	{Jet_pt, nJet, nMuon}	ValueError: cannot broadcast nested list
uproot.graphed, mapped	{Jet_pt, nJet, nMuon}	ValueError: cannot broadcast nested list (head and write)
eager ak.zip on the read arrays	—	ValueError: cannot broadcast nested list

So Muon_pt is omitted by design — dask-awkward projects it away too — and what fails is awkward's broadcast of jets against muons (different per-event multiplicities), identically on every path. graphed does not read a column dask would not read, and does not fail where dask succeeds.

The well-formed spelling, zip({...}, depth_limit=1)[["jet"]], projects to {Jet_pt, nJet, nMuon} and evaluates on graphed_head and graphed_write (test_unread_branches_are_placeholders_the_graph_can_pass_through). Side note from the same probe: through this mapping uproot.dask reads only {Jet_pt, nJet} for that spelling and its compute() fails with TypeError: PlaceholderArray supports only trivial slices, not int — the muon list needs its nMuon offsets even when its data is never touched; graphed declares that buffer as an offsets need and reads it.

[lgray]

🤖 AI text below 🤖

P2: the default process writer cannot ship common behaviors … ROOT writing needs the same seam or should avoid serializing the backend instance.

Fixed with the same seam as the Parquet writer. The worker receives the backend as a class or a "module:attr" import reference (the aggregate_plan convention), default = the session backend's class; a behavior travels by reference (test_process_pool_with_a_behavior_backend_by_import_ref, which also pins that a lambda-bearing backend instance still cannot pickle — that is inherent, and no longer the default).

[lgray]

🤖 AI text below 🤖

P2: external evaluators are dropped by graphed_head and graphed_write … aggregate_plan already wires these through _external_evaluators.

Fixed. Both carry the session's evaluators, through the now-public graphed.aggregate.external_evaluators (graphed-org/graphed#38, which also wired them into graphed's own writers, which had the same gap). test_externals_reach_the_write_workers; graphed_head(graphed.apply(f, g.px1), 2) evaluates.

[lgray]

🤖 AI text below 🤖

P2: already-open TTrees are not partition-compatible … Normalize it to a reusable file/tree reference or special-case open objects in open_tree.

Fixed by normalising: an open TTree becomes its (file_path, object_path) pair and workers reopen it by path (test_an_open_tree_is_reopened_by_path_in_the_workers).

[lgray]

@ariostas OK - bot responses done and everything is addressed. Re-review or let me know what to do otherwise! Thanks!