feat: Add RNTuple row extension and field addition via update mode - #1687
feat: Add RNTuple row extension and field addition via update mode#1687Yokubas wants to merge 75 commits into
Conversation
…s/uproot5 into Yokubas/rntuple-update-pr
…structing full header
…s/uproot5 into Yokubas/rntuple-update-pr
…s/uproot5 into Yokubas/rntuple-update-pr
Codecov Report❌ Patch coverage is
❌ Your patch check has failed because the patch coverage (94.71%) is below the target coverage (98.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files
|
…s/uproot5 into Yokubas/rntuple-update-pr
…s/uproot5 into Yokubas/rntuple-update-pr
…s/uproot5 into Yokubas/rntuple-update-pr
| if len(ple.pagelinklist) > 1: | ||
| raise ValueError( | ||
| f"add_fields does not yet support RNTuples with multiple clusters per cluster group " | ||
| f"(cluster group {cg_idx} has {len(ple.pagelinklist)} clusters). " | ||
| f"This is a known limitation that will be fixed in a future version." | ||
| ) |
There was a problem hiding this comment.
We should implement this since I think it should be pretty straightforward. But let me know if it's not so easy and we can leave it for later.
There was a problem hiding this comment.
I implemented multi-cluster support. Instead of raising an error, add_fields now loops over all clusters in each cluster group and writes one zero-filled page per new field per cluster, sized to that cluster's entry span. Tested with 3 cluster groups and it works correctly. Note that multiple clusters per cluster group (ROOT behavior) can't be easily tested since ROOT-written files have split encoding which we reject — but the loop handles that case correctly too.
| key = self._file.root_directory._cascading.data.get_key( | ||
| self._path[-1], 1 | ||
| ) | ||
| reloaded = self._file.root_directory._load_existing_ntuple(key) |
There was a problem hiding this comment.
Maybe it's better for add_fields to update the necessary things instead of reloading the entire ntuple
There was a problem hiding this comment.
Reduced the file reload — _existing_field_records and _column_counts are now updated directly in memory. Still need to read _existing_footer, _existing_page_list_envelopes, and _header._akform from file because the writable footer uses cluster_group_record_frames while the next add_fields call needs cluster_group_records from the read-only footer. Could avoid this by adding a cluster_group_records property to the writable footer that mirrors cluster_group_record_frames — would that be the preferred approach?
…s/uproot5 into Yokubas/rntuple-update-pr
…s/uproot5 into Yokubas/rntuple-update-pr
| ) | ||
| footer.extension_column_record_frames.append(new_col) | ||
|
|
||
| new_data = numpy.zeros(num_entries, dtype=numpy.dtype(ak_primitive)) |
There was a problem hiding this comment.
Sorry, I missed this part on my first review. We actually, don't want to explicitly backfill the new fields with zeros. This is taken care by attaching a deferred column to the new field. So the existing cluster groups stay as they were, with no new columns for the new fields, and new cluster groups will start to contain the new column. So you'll have to modify NTuple_Column_Description so that it can also take a first-element index (see https://github.com/root-project/root/blob/master/tree/ntuple/doc/BinaryFormatSpecification.md#column-description).
So now it makes sense why there were some issues with cluster groups with multiple clusters. This way, since you're leaving old cluster groups intact and Uproot (currently) only writes cluster groups with a single cluster then you don't have to worry about multiple clusters in a cluster group.
There was a problem hiding this comment.
Implemented the deferred column approach — NTuple_Column_Description now accepts first_element_index, and when it's > 0 the DEFERRED flag is set and the value is serialized. add_fields no longer backfills zeros — old cluster groups stay completely untouched, and the reader pads zeros automatically via the deferred column mechanism. Also fixed the reload path to preserve first_element_index when copying extension column records. 32 tests still pass.
…s/uproot5 into Yokubas/rntuple-update-pr
Summary
Implements in-place modification of existing RNTuples:
f["name"].extend({"x": array1, "y": array2})— append new rows to existing RNTuplef["name"].add_fields({"z": np.int32, "w": np.float32, ...})— add one or more new fields back-filled with zerosf["name"].add_fields({"particle.phi": np.float32, "particle.eta": np.float64, ...})— add subfields to existing untyped structsf["name"].extend({"x": array1, "z": array2}, accept_new_fields=True)— auto-add new fields and extendHow it works
first_element_index = num_entries, marking where new data starts. Old cluster groups are left completely untouched. The reader automatically zero-pads entries beforefirst_element_index. Subsequentextendcalls write new cluster groups that include the new column.Tests
32 tests in
tests/test_1687_rntuple_update.pycovering:Blocking issues addressed (from review)
ValueErrorwith encoding mismatch instead of writing garbage datanum_entrieselementsKnown limitations