From 9e401400b7ae18b2a5b98fdbcc6c8384feaf823f Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:18:56 +0000 Subject: [PATCH 1/9] Give page descriptions their column, field and global cluster ID RNTuple.get_extended_page_descriptions(): - returns one entry per physical column in every cluster, so a position is the column ID, as the spec guarantees (Page Locations: the outer items follow column order). With the old default includeSuppressed=False, suppressed columns were dropped and positions shifted. The parameter is removed; a suppressed column is an entry with no pages. The zip over columns is strict, as suggested on #49. - gives each page its RNTuple-wide cluster ID. Cluster IDs continue across cluster groups (Page List Envelope); the new RNTuple.firstClusterIDs is the running sum of the footer's fNClusters. Each page list's cluster count is checked against its group. - gives each page its column ID, field ID, FieldDescription and qualified field path. SchemaDescription.field_path(field_id) walks fParentFieldID up to the top-level field (which names itself as parent) and joins the names with "." (forbidden in names, so unambiguous). It returns bytes, like the names, and raises on an out-of-range ID or a cycle. The schema is built once per call; the schemaDescription property itself is unchanged. Fixes #49 Fixes #116 Assisted-by: claude-code:claude-opus-5-5 --- src/rootfilespec/rntuple/RNTuple.py | 141 +++++++++++++++++++----- tests/test_rntuple_hardcoded.py | 42 +++++++ tests/test_rntuple_page_descriptions.py | 121 ++++++++++++++++++++ 3 files changed, 274 insertions(+), 30 deletions(-) create mode 100644 tests/test_rntuple_page_descriptions.py diff --git a/src/rootfilespec/rntuple/RNTuple.py b/src/rootfilespec/rntuple/RNTuple.py index a5c9aa6..42bf013 100644 --- a/src/rootfilespec/rntuple/RNTuple.py +++ b/src/rootfilespec/rntuple/RNTuple.py @@ -79,6 +79,35 @@ def from_envelopes( extraTypeInformations, ) + def field_path(self, field_id: int) -> bytes: + """The qualified name of a field: its ancestors' names and its own, joined by "." + + A field ID is the field's position in the combined list (header first, + then the footer's schema extension). A top-level field names itself as + its parent, which ends the walk. "." is forbidden in field names (spec, + *Naming specification*), so the result is unambiguous. Names are bytes, + as stored. + """ + fields = self.fieldDescriptions + names: list[bytes] = [] + seen: set[int] = set() + fid = field_id + while True: + if not 0 <= fid < len(fields): + msg = f"Field {field_id}: parent chain reaches field ID {fid}, of {len(fields)} fields" + raise ValueError(msg) + if fid in seen: + msg = ( + f"Field {field_id}: the parent chain has a cycle at field ID {fid}" + ) + raise ValueError(msg) + seen.add(fid) + field = fields[fid] + names.append(field.fFieldName) + if field.fParentFieldID == fid: + return b".".join(reversed(names)) + fid = field.fParentFieldID + @dataclasses.dataclass class InterpretablePage: @@ -92,6 +121,16 @@ class InterpretablePage: """The uncompressed size of the page, in bytes.""" columnType: ColumnType """The type of the column this page belongs to, e.g. kInt32, kFloat64, etc.""" + clusterID: int + """The RNTuple-wide ID of the page's cluster, continuous across cluster groups.""" + columnID: int + """The ID of the page's physical column (its position in the combined column list).""" + fieldID: int + """The ID of the field the column belongs to (``ColumnDescription.fFieldID``).""" + fieldDescription: FieldDescription + """The description of that field.""" + fieldPath: bytes + """The field's qualified name (see ``SchemaDescription.field_path``).""" @dataclasses.dataclass @@ -173,40 +212,82 @@ def streamer_infos(self) -> dict[bytes, TStreamerInfo]: infos[info.fName] = info return infos - # can provide helpers to get page descriptions with different filters, columns/rows/etc. + @property + def firstClusterIDs(self) -> list[int]: + """The ID of the first cluster of each cluster group, in footer order + + Cluster IDs continue across cluster groups: a page list's clusters start + at the number of clusters in all earlier groups (spec, *Page List + Envelope*). ``pagelistEnvelopes`` are in the same order as the groups. + """ + first, out = 0, [] + for group in self.footerEnvelope.clusterGroups: + out.append(first) + first += group.fNClusters + return out + def get_extended_page_descriptions( self, - includeSuppressed: bool = False, ) -> list[list[list[list[InterpretablePage]]]]: - """Fetches all pages from the RNTuple, organized by cluster group, column, and page. + """All page descriptions, by page-list envelope, cluster, column and page. - Args: - includeSuppressed (bool): If False, skip suppressed columns. + The column list of each cluster has one entry per physical column, in + column-ID order, as the spec guarantees: a suppressed column is there + too, with no pages. Each page carries its global cluster ID, its column + ID and its field. """ - envelopePages: list[list[list[list[InterpretablePage]]]] = [ - [ - [ - [ - InterpretablePage( - pageDescription=page_description, - uncompressedSize=ceil( - abs(page_description.fNElements) - * column_description.fBitsOnStorage - / 8 - ), # Convert bits to bytes - columnType=column_description.fColumnType, - ) - for page_description in pagelist - ] - for pagelist, column_description in zip( - columnlist, - self.schemaDescription.columnDescriptions, - strict=False, + schema = self.schemaDescription + columns = schema.columnDescriptions + paths: dict[int, bytes] = {} + if len(self.pagelistEnvelopes) != len(self.footerEnvelope.clusterGroups): + msg = ( + f"{len(self.pagelistEnvelopes)} page lists for " + f"{len(self.footerEnvelope.clusterGroups)} cluster groups" + ) + raise ValueError(msg) + envelopePages: list[list[list[list[InterpretablePage]]]] = [] + for pagelistEnvelope, group, firstClusterID in zip( + self.pagelistEnvelopes, + self.footerEnvelope.clusterGroups, + self.firstClusterIDs, + strict=True, + ): + if len(pagelistEnvelope.pageLocations) != group.fNClusters: + msg = ( + f"Page list of cluster group starting at cluster {firstClusterID} " + f"has {len(pagelistEnvelope.pageLocations)} clusters, " + f"the group says {group.fNClusters}" + ) + raise ValueError(msg) + clusters: list[list[list[InterpretablePage]]] = [] + for index, columnlist in enumerate(pagelistEnvelope.pageLocations): + clusterID = firstClusterID + index + clusterPages: list[list[InterpretablePage]] = [] + for columnID, (pagelist, column) in enumerate( + zip(columnlist, columns, strict=True) + ): + fieldID = column.fFieldID + if fieldID not in paths: + paths[fieldID] = schema.field_path(fieldID) + clusterPages.append( + [ + InterpretablePage( + pageDescription=page_description, + uncompressedSize=ceil( + abs(page_description.fNElements) + * column.fBitsOnStorage + / 8 + ), # Convert bits to bytes + columnType=column.fColumnType, + clusterID=clusterID, + columnID=columnID, + fieldID=fieldID, + fieldDescription=schema.fieldDescriptions[fieldID], + fieldPath=paths[fieldID], + ) + for page_description in pagelist + ] ) - if includeSuppressed or pagelist.elementoffset >= 0 - ] - for columnlist in pagelistEnvelope.pageLocations - ] - for pagelistEnvelope in self.pagelistEnvelopes - ] + clusters.append(clusterPages) + envelopePages.append(clusters) return envelopePages diff --git a/tests/test_rntuple_hardcoded.py b/tests/test_rntuple_hardcoded.py index 1b743af..caac81f 100644 --- a/tests/test_rntuple_hardcoded.py +++ b/tests/test_rntuple_hardcoded.py @@ -350,6 +350,13 @@ def test_read_contributors(): ), uncompressedSize=176, columnType=ColumnType.kIndex64, + clusterID=0, + columnID=0, + fieldID=0, + fieldDescription=rntuple.schemaDescription.fieldDescriptions[ + 0 + ], + fieldPath=b"firstName", ) ], [ @@ -360,6 +367,13 @@ def test_read_contributors(): ), uncompressedSize=178, columnType=ColumnType.kChar, + clusterID=0, + columnID=1, + fieldID=0, + fieldDescription=rntuple.schemaDescription.fieldDescriptions[ + 0 + ], + fieldPath=b"firstName", ) ], [ @@ -370,6 +384,13 @@ def test_read_contributors(): ), uncompressedSize=176, columnType=ColumnType.kIndex64, + clusterID=0, + columnID=2, + fieldID=1, + fieldDescription=rntuple.schemaDescription.fieldDescriptions[ + 1 + ], + fieldPath=b"lastName", ) ], [ @@ -380,6 +401,13 @@ def test_read_contributors(): ), uncompressedSize=193, columnType=ColumnType.kChar, + clusterID=0, + columnID=3, + fieldID=1, + fieldDescription=rntuple.schemaDescription.fieldDescriptions[ + 1 + ], + fieldPath=b"lastName", ) ], ] @@ -598,6 +626,13 @@ def fetch_from_locator(loc): ), uncompressedSize=400, columnType=ColumnType.kSplitReal32, + clusterID=0, + columnID=0, + fieldID=0, + fieldDescription=rntuple_a.schemaDescription.fieldDescriptions[ + 0 + ], + fieldPath=b"f", ) ] ] @@ -791,6 +826,13 @@ def fetch_from_locator(loc): ), uncompressedSize=400, columnType=ColumnType.kSplitInt32, + clusterID=0, + columnID=0, + fieldID=0, + fieldDescription=rntuple_b.schemaDescription.fieldDescriptions[ + 0 + ], + fieldPath=b"g", ) ] ] diff --git a/tests/test_rntuple_page_descriptions.py b/tests/test_rntuple_page_descriptions.py new file mode 100644 index 0000000..dd44315 --- /dev/null +++ b/tests/test_rntuple_page_descriptions.py @@ -0,0 +1,121 @@ +import dataclasses +from pathlib import Path + +import pytest +from skhep_testdata import data_path # type: ignore[import-not-found] + +from rootfilespec.reader import open_path +from rootfilespec.rntuple.RNTuple import RNTuple, SchemaDescription +from rootfilespec.rntuple.schema import ColumnType + +DATA = Path(__file__).parent.parent / "reference" / "root-io-spec" / "data" / "rntuple" + + +def _load(path: str | Path) -> RNTuple: + with open_path(path) as reader: + keylist = reader.keylist() + (name,) = [ + n for n in keylist if keylist[n].fClassName.fString == b"ROOT::RNTuple" + ] + return RNTuple.from_anchor(reader.fetch(keylist[name]), reader.fetch.buffer) + + +def test_suppressed_columns_keep_their_position(): + """A column's position is its column ID, suppressed or not + + In this file the field "real" has two representations, Real32 (column 0) + and Real16 (column 1), and each cluster suppresses one of them. Before, + cluster 1's only entry was column 1, at position 0. + """ + rntuple = _load(data_path("test_multiple_representations_rntuple_v1-0-0-0.root")) + (envelope,) = rntuple.get_extended_page_descriptions() + assert [[len(column) for column in cluster] for cluster in envelope] == [ + [1, 0], + [0, 1], + [1, 0], + ] + (page,) = envelope[1][1] + assert page.columnID == 1 + assert page.columnType == ColumnType.kReal16 + suppressed = rntuple.pagelistEnvelopes[0].pageLocations[1][0] + assert suppressed.elementoffset < 0 + + +def test_cluster_ids_continue_across_cluster_groups(): + """Issue #116: 3 cluster groups of 5, 4 and 3 clusters are clusters 0 to 11""" + rntuple = _load(data_path("test_multiple_cluster_groups_rntuple_v1-0-0-0.root")) + assert rntuple.firstClusterIDs == [0, 5, 9] + ids = [ + {page.clusterID for column in cluster for page in column} + for envelope in rntuple.get_extended_page_descriptions() + for cluster in envelope + ] + assert ids == [{i} for i in range(12)] + first_entries = [ + summary.fFirstEntryNumber + for pagelist in rntuple.pagelistEnvelopes + for summary in pagelist.clusterSummaries + ] + assert first_entries == [0, 100, 200, 300, 400, 450, 500, 600, 700, 750, 800, 900] + + +@pytest.mark.skipif(not DATA.exists(), reason="reference/root-io-spec not checked out") +def test_field_paths_user_class(): + """Issue #49: qualified field names, as root-io-spec's case for this file names + them (gen/cases/rntuple/user-class/case.toml)""" + schema = _load(DATA / "user-class.root").schemaDescription + names = {i: f.fFieldName for i, f in enumerate(schema.fieldDescriptions)} + paths = {i: schema.field_path(i) for i in names} + assert paths[0] == b"fHit" + assert paths[1] == b"fHit.:_0" + assert paths[2] == b"fHit.:_0.fBaseId" + assert paths[13] == b"fHits._0.:_0" + assert paths[24] == b"fFlavour._0" + assert paths[26] == b"fCharge._0" + for i, path in paths.items(): + assert path.rsplit(b".", 1)[-1] == names[i] + + +def test_field_paths_schema_extension(): + """Field IDs continue from the header into the footer's schema extension""" + rntuple = _load(data_path("test_extension_columns_rntuple_v1-0-0-0.root")) + assert len(rntuple.headerEnvelope.fieldDescriptions) == 1 + schema = rntuple.schemaDescription + assert [schema.field_path(i) for i in range(4)] == [ + b"int_field", + b"float_field", + b"intvec_field", + b"intvec_field._0", + ] + + +@pytest.mark.skipif(not DATA.exists(), reason="reference/root-io-spec not checked out") +@pytest.mark.parametrize("name", sorted(p.name for p in DATA.glob("*.root"))) +def test_every_page_knows_its_column_and_field(name: str): + rntuple = _load(DATA / name) + schema = rntuple.schemaDescription + for envelope in rntuple.get_extended_page_descriptions(): + for cluster in envelope: + assert len(cluster) == len(schema.columnDescriptions) + for columnID, column in enumerate(cluster): + for page in column: + assert page.columnID == columnID + field_id = schema.columnDescriptions[columnID].fFieldID + assert page.fieldID == field_id + assert page.fieldDescription == schema.fieldDescriptions[field_id] + assert page.fieldPath == schema.field_path(field_id) + + +def test_field_path_errors(): + schema = _load( + data_path("test_extension_columns_rntuple_v1-0-0-0.root") + ).schemaDescription + with pytest.raises(ValueError, match="parent chain reaches field ID 9"): + schema.field_path(9) + fields = list(schema.fieldDescriptions) + # _0 -> intvec_field -> _0: a cycle + fields[2] = dataclasses.replace(fields[2], fParentFieldID=3) + broken = dataclasses.replace(schema, fieldDescriptions=fields) + assert isinstance(broken, SchemaDescription) + with pytest.raises(ValueError, match="cycle"): + broken.field_path(3) From 19e6bc4aaae7ccae8f63ce890b2f714aeb385c11 Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:18:36 +0000 Subject: [PATCH 2/9] Rename tests/test_rntuple_page_descriptions.py to tests/test_extended_page_descriptions.py Name the module after what it tests, so it does not read as the home of a broad set of tests (review on #125). Assisted-by: claude-code:claude-opus-5-5 --- ...le_page_descriptions.py => test_extended_page_descriptions.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{test_rntuple_page_descriptions.py => test_extended_page_descriptions.py} (100%) diff --git a/tests/test_rntuple_page_descriptions.py b/tests/test_extended_page_descriptions.py similarity index 100% rename from tests/test_rntuple_page_descriptions.py rename to tests/test_extended_page_descriptions.py From 73d652408548f74093201265112ade950e3e0d28 Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:53:13 +0000 Subject: [PATCH 3/9] Compare key class names as bytes, after the rest of #68 in #127 Assisted-by: claude-code:claude-opus-5-5 --- tests/test_extended_page_descriptions.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_extended_page_descriptions.py b/tests/test_extended_page_descriptions.py index dd44315..535e4c4 100644 --- a/tests/test_extended_page_descriptions.py +++ b/tests/test_extended_page_descriptions.py @@ -14,9 +14,7 @@ def _load(path: str | Path) -> RNTuple: with open_path(path) as reader: keylist = reader.keylist() - (name,) = [ - n for n in keylist if keylist[n].fClassName.fString == b"ROOT::RNTuple" - ] + (name,) = [n for n in keylist if keylist[n].fClassName == b"ROOT::RNTuple"] return RNTuple.from_anchor(reader.fetch(keylist[name]), reader.fetch.buffer) From 0e3fa2948be54d7ebc52841b5e41c52259f33964 Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Wed, 30 Sep 2026 19:47:40 -0500 Subject: [PATCH 4/9] Accept page lists that predate a model extension A cluster committed before the model was extended lists only the columns that existed then, so its page list can end before the schema extension's columns do (the spec fixes the columns' order, not their number). ROOT fills those in when reading (AddExtendedColumnRanges, RNTupleDescriptor.cxx:920 at 6.40.04). The strict zip raised on scikit-hep-testdata's test_extension_columns_rntuple_v1-0-0-0.root, which read before this PR: its first cluster lists 2 of 4 columns. get_extended_page_descriptions now gives such trailing columns empty entries, so a position is still the column ID. A cluster listing more columns than the schema, or fewer than the header's, still raises. uncompressedSize uses RPageDescription.n_elements (#128). Tests: that file's clusters, a page list missing a header column, and the per-page column/field check on three scikit-hep-testdata files as well as root-io-spec's. Assisted-by: claude-code:claude-opus-5-5 --- src/rootfilespec/rntuple/RNTuple.py | 22 ++++++++- tests/test_extended_page_descriptions.py | 57 ++++++++++++++++++++++-- 2 files changed, 73 insertions(+), 6 deletions(-) diff --git a/src/rootfilespec/rntuple/RNTuple.py b/src/rootfilespec/rntuple/RNTuple.py index 42bf013..4761a3e 100644 --- a/src/rootfilespec/rntuple/RNTuple.py +++ b/src/rootfilespec/rntuple/RNTuple.py @@ -235,9 +235,16 @@ def get_extended_page_descriptions( column-ID order, as the spec guarantees: a suppressed column is there too, with no pages. Each page carries its global cluster ID, its column ID and its field. + + A cluster committed before the model was extended lists only the + columns that existed then, so its page list can end before the schema + extension's columns do. Those trailing columns get empty entries, as + ROOT synthesizes them when reading (``AddExtendedColumnRanges``, + ``RNTupleDescriptor.cxx:920`` at 6.40.04). Any other mismatch raises. """ schema = self.schemaDescription columns = schema.columnDescriptions + nHeaderColumns = len(self.headerEnvelope.columnDescriptions.items) paths: dict[int, bytes] = {} if len(self.pagelistEnvelopes) != len(self.footerEnvelope.clusterGroups): msg = ( @@ -262,9 +269,20 @@ def get_extended_page_descriptions( clusters: list[list[list[InterpretablePage]]] = [] for index, columnlist in enumerate(pagelistEnvelope.pageLocations): clusterID = firstClusterID + index + if not nHeaderColumns <= len(columnlist) <= len(columns): + msg = ( + f"Cluster {clusterID} lists {len(columnlist)} columns; the " + f"schema has {nHeaderColumns} in the header and " + f"{len(columns)} in all" + ) + raise ValueError(msg) + # Columns added by a later model extension, which the cluster predates + missing: list[list[RPageDescription]] = [ + [] for _ in range(len(columns) - len(columnlist)) + ] clusterPages: list[list[InterpretablePage]] = [] for columnID, (pagelist, column) in enumerate( - zip(columnlist, columns, strict=True) + zip([*columnlist, *missing], columns, strict=True) ): fieldID = column.fFieldID if fieldID not in paths: @@ -274,7 +292,7 @@ def get_extended_page_descriptions( InterpretablePage( pageDescription=page_description, uncompressedSize=ceil( - abs(page_description.fNElements) + page_description.n_elements * column.fBitsOnStorage / 8 ), # Convert bits to bytes diff --git a/tests/test_extended_page_descriptions.py b/tests/test_extended_page_descriptions.py index 535e4c4..da4d357 100644 --- a/tests/test_extended_page_descriptions.py +++ b/tests/test_extended_page_descriptions.py @@ -87,10 +87,59 @@ def test_field_paths_schema_extension(): ] -@pytest.mark.skipif(not DATA.exists(), reason="reference/root-io-spec not checked out") -@pytest.mark.parametrize("name", sorted(p.name for p in DATA.glob("*.root"))) -def test_every_page_knows_its_column_and_field(name: str): - rntuple = _load(DATA / name) +# scikit-hep-testdata files whose page lists differ from the root-io-spec ones: +# suppressed columns, several cluster groups, and a model extended after the +# first cluster was committed +SKHEP_FILES = [ + "test_multiple_representations_rntuple_v1-0-0-0.root", + "test_multiple_cluster_groups_rntuple_v1-0-0-0.root", + "test_extension_columns_rntuple_v1-0-0-0.root", +] + + +def test_columns_added_by_model_extension(): + """A cluster committed before the model was extended lists only the columns + that existed then: here 2 of the 4. The two later ones get empty entries, + as ROOT synthesizes them (AddExtendedColumnRanges), so a position is still + the column ID.""" + rntuple = _load(data_path("test_extension_columns_rntuple_v1-0-0-0.root")) + (pagelist,) = rntuple.pagelistEnvelopes + assert [len(columns) for columns in pagelist.pageLocations] == [2, 4, 4, 4] + (envelope,) = rntuple.get_extended_page_descriptions() + assert [[len(column) for column in cluster] for cluster in envelope] == [ + [2, 1, 0, 0], + [1, 1, 1, 1], + [1, 1, 1, 1], + [1, 1, 1, 1], + ] + + +def test_page_list_with_a_missing_header_column_raises(): + """Only columns of the schema extension may be missing from a page list""" + rntuple = _load(data_path("test_extension_columns_rntuple_v1-0-0-0.root")) + (pagelist,) = rntuple.pagelistEnvelopes + pagelist.pageLocations.items[0].items = [] + with pytest.raises(ValueError, match="Cluster 0 lists 0 columns"): + rntuple.get_extended_page_descriptions() + + +@pytest.mark.parametrize( + "path", + [ + *( + pytest.param(DATA / name, id=name) + for name in sorted(p.name for p in DATA.glob("*.root")) + ), + *(pytest.param(name, id=name) for name in SKHEP_FILES), + ], +) +def test_every_page_knows_its_column_and_field(path: Path | str): + if isinstance(path, Path): + if not path.exists(): + pytest.skip("reference/root-io-spec not checked out") + rntuple = _load(path) + else: + rntuple = _load(data_path(path)) schema = rntuple.schemaDescription for envelope in rntuple.get_extended_page_descriptions(): for cluster in envelope: From 37513a210727362a82c45f26fe422144c3c983ba Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:34:43 -0500 Subject: [PATCH 5/9] Give each column in a cluster its element range, not just its pages get_extended_page_descriptions() now returns one InterpretableColumn per physical column per cluster, so that the view is enough on its own to read a column. Before, a column was a bare list of pages, and an empty list could mean a suppressed column, a column with no pages, or a column the cluster predates; the element offset was not in the view at all. The ranges are the ones ROOT's reader builds (CommitSuppressedColumnRanges and AddExtendedColumnRanges, tree/ntuple/src/RNTupleDescriptor.cxx:880 and :920 at 6.40.04), which is what the spec asks of a reader: - a suppressed column takes the range of the corresponding column of the field's active representation (spec, Suppressed Columns); - a deferred column covers every cluster from element 0, the elements before its first stored one being zeros with no page (Column Description): nZeroElements counts them; - a column the cluster predates keeps pageLocations=None. New checks, each with a test: a cluster cannot list fewer columns than an earlier one, since only a model extension adds columns; a page list needs a cluster summary per cluster; a suppressed column has no pages; every field has an active representation; a deferred column's pages start where its range says; and a deferred column may not be inside a collection or variant, whose element count the format does not record. The tests now also pin each column's field path and type against root-io-spec's user-class case, and check on every fixture that a column's clusters follow each other with no gap. Assisted-by: claude-code:claude-opus-5-5 --- src/rootfilespec/rntuple/RNTuple.py | 385 ++++++++++++++++++----- tests/test_extended_page_descriptions.py | 296 +++++++++++++---- tests/test_rntuple_hardcoded.py | 15 +- 3 files changed, 560 insertions(+), 136 deletions(-) diff --git a/src/rootfilespec/rntuple/RNTuple.py b/src/rootfilespec/rntuple/RNTuple.py index 4761a3e..05dd26d 100644 --- a/src/rootfilespec/rntuple/RNTuple.py +++ b/src/rootfilespec/rntuple/RNTuple.py @@ -10,8 +10,8 @@ from rootfilespec.rntuple.envelope import RFeatureFlags from rootfilespec.rntuple.footer import FooterEnvelope from rootfilespec.rntuple.header import HeaderEnvelope -from rootfilespec.rntuple.pagelist import PageListEnvelope -from rootfilespec.rntuple.pagelocations import RPageDescription +from rootfilespec.rntuple.pagelist import ClusterSummary, PageListEnvelope +from rootfilespec.rntuple.pagelocations import PageLocations, RPageDescription from rootfilespec.rntuple.schema import ( AliasColumnDescription, ColumnDescription, @@ -79,34 +79,38 @@ def from_envelopes( extraTypeInformations, ) - def field_path(self, field_id: int) -> bytes: - """The qualified name of a field: its ancestors' names and its own, joined by "." + def field_chain(self, field_id: int) -> list[int]: + """The IDs of a field and of its ancestors, from the field up to its top-level field A field ID is the field's position in the combined list (header first, then the footer's schema extension). A top-level field names itself as - its parent, which ends the walk. "." is forbidden in field names (spec, - *Naming specification*), so the result is unambiguous. Names are bytes, - as stored. + its parent, which ends the walk. Raises on an out-of-range ID or a cycle. """ fields = self.fieldDescriptions - names: list[bytes] = [] - seen: set[int] = set() + chain: list[int] = [] fid = field_id while True: if not 0 <= fid < len(fields): msg = f"Field {field_id}: parent chain reaches field ID {fid}, of {len(fields)} fields" raise ValueError(msg) - if fid in seen: + if fid in chain: msg = ( f"Field {field_id}: the parent chain has a cycle at field ID {fid}" ) raise ValueError(msg) - seen.add(fid) - field = fields[fid] - names.append(field.fFieldName) - if field.fParentFieldID == fid: - return b".".join(reversed(names)) - fid = field.fParentFieldID + chain.append(fid) + if fields[fid].fParentFieldID == fid: + return chain + fid = fields[fid].fParentFieldID + + def field_path(self, field_id: int) -> bytes: + """The qualified name of a field: its ancestors' names and its own, joined by "." + + "." is forbidden in field names (spec, *Naming specification*), so the + result is unambiguous. Names are bytes, as stored. + """ + chain = self.field_chain(field_id) + return b".".join(self.fieldDescriptions[fid].fFieldName for fid in chain[::-1]) @dataclasses.dataclass @@ -133,6 +137,57 @@ class InterpretablePage: """The field's qualified name (see ``SchemaDescription.field_path``).""" +@dataclasses.dataclass +class InterpretableColumn: + """One physical column in one cluster: its element range and its pages + + Enough on its own to read the column's elements in the cluster. The element + range is the one ROOT's reader builds (``CommitSuppressedColumnRanges`` and + ``AddExtendedColumnRanges``, ``tree/ntuple/src/RNTupleDescriptor.cxx:880`` + and ``:920`` at 6.40.04), which is what the spec asks of a reader + (*Suppressed Columns*; *Column Description* for deferred columns). + """ + + clusterID: int + """The RNTuple-wide ID of the cluster, continuous across cluster groups.""" + columnID: int + """The ID of the physical column (its position in the combined column list).""" + columnDescription: ColumnDescription + """The description of the column.""" + fieldID: int + """The ID of the field the column belongs to (``ColumnDescription.fFieldID``).""" + fieldDescription: FieldDescription + """The description of that field.""" + fieldPath: bytes + """The field's qualified name (see ``SchemaDescription.field_path``).""" + pageLocations: PageLocations[RPageDescription] | None + """The column's entry in the cluster's page list, as stored. + + ``None`` if the page list has no entry for the column: a cluster committed + before the model was extended lists only the columns that existed then + (ROOT's ``SerializePageList``, ``tree/ntuple/src/RNTupleSerialize.cxx:1693`` + at 6.40.04).""" + suppressed: bool + """Whether the column is a secondary representation, inactive in this cluster. + + A suppressed column has no pages, and its element range is that of the + corresponding column of the field's active representation. A column that + the cluster predates is suppressed if it is deferred and suppressed + (negative ``fFirstElementIndex``).""" + firstElementIndex: int + """The index, within the column, of its first element in this cluster.""" + nElements: int + """The number of elements of the column in this cluster.""" + nZeroElements: int + """The number of leading elements that have no page on disk and read as zeros. + + Non-zero only for a deferred column, in the clusters up to its first stored + element: a reader yields zero bytes for them (spec, *Column Description*). + The pages hold the remaining ``nElements - nZeroElements``.""" + pages: list[InterpretablePage] + """The column's pages on disk in this cluster, in order.""" + + @dataclasses.dataclass class RNTuple: """A class representing an RNTuple.""" @@ -228,84 +283,254 @@ def firstClusterIDs(self) -> list[int]: def get_extended_page_descriptions( self, - ) -> list[list[list[list[InterpretablePage]]]]: - """All page descriptions, by page-list envelope, cluster, column and page. - - The column list of each cluster has one entry per physical column, in - column-ID order, as the spec guarantees: a suppressed column is there - too, with no pages. Each page carries its global cluster ID, its column - ID and its field. - - A cluster committed before the model was extended lists only the - columns that existed then, so its page list can end before the schema - extension's columns do. Those trailing columns get empty entries, as - ROOT synthesizes them when reading (``AddExtendedColumnRanges``, - ``RNTupleDescriptor.cxx:920`` at 6.40.04). Any other mismatch raises. + ) -> list[list[list[InterpretableColumn]]]: + """Every column of every cluster, by page-list envelope and cluster. + + Each cluster has one entry per physical column, in column-ID order, so a + position is the column ID (spec, *Page Locations*). Each entry is enough + on its own to read the column in that cluster; see + ``InterpretableColumn``. + + Raises ``ValueError`` where the page lists contradict the footer or the + schema. """ schema = self.schemaDescription - columns = schema.columnDescriptions + nColumns = len(schema.columnDescriptions) nHeaderColumns = len(self.headerEnvelope.columnDescriptions.items) - paths: dict[int, bytes] = {} - if len(self.pagelistEnvelopes) != len(self.footerEnvelope.clusterGroups): - msg = ( - f"{len(self.pagelistEnvelopes)} page lists for " - f"{len(self.footerEnvelope.clusterGroups)} cluster groups" - ) + groups = self.footerEnvelope.clusterGroups + if len(self.pagelistEnvelopes) != len(groups): + msg = f"{len(self.pagelistEnvelopes)} page lists for {len(groups)} cluster groups" raise ValueError(msg) - envelopePages: list[list[list[list[InterpretablePage]]]] = [] + context = _ColumnContext(schema) + envelopes: list[list[list[InterpretableColumn]]] = [] + # The cluster that lists the most columns so far, and how many: only a + # model extension adds columns, so no later cluster lists fewer + widestClusterID, widest = 0, 0 for pagelistEnvelope, group, firstClusterID in zip( - self.pagelistEnvelopes, - self.footerEnvelope.clusterGroups, - self.firstClusterIDs, - strict=True, + self.pagelistEnvelopes, groups, self.firstClusterIDs, strict=True ): - if len(pagelistEnvelope.pageLocations) != group.fNClusters: + pageLocations = pagelistEnvelope.pageLocations.items + summaries = pagelistEnvelope.clusterSummaries.items + if not len(pageLocations) == len(summaries) == group.fNClusters: msg = ( f"Page list of cluster group starting at cluster {firstClusterID} " - f"has {len(pagelistEnvelope.pageLocations)} clusters, " - f"the group says {group.fNClusters}" + f"has {len(summaries)} cluster summaries and page locations for " + f"{len(pageLocations)} clusters, the group says {group.fNClusters}" ) raise ValueError(msg) - clusters: list[list[list[InterpretablePage]]] = [] - for index, columnlist in enumerate(pagelistEnvelope.pageLocations): + clusters: list[list[InterpretableColumn]] = [] + for index, (columnlist, summary) in enumerate( + zip(pageLocations, summaries, strict=True) + ): clusterID = firstClusterID + index - if not nHeaderColumns <= len(columnlist) <= len(columns): + listed = columnlist.items + if not nHeaderColumns <= len(listed) <= nColumns: msg = ( - f"Cluster {clusterID} lists {len(columnlist)} columns; the " + f"Cluster {clusterID} lists {len(listed)} columns; the " f"schema has {nHeaderColumns} in the header and " - f"{len(columns)} in all" + f"{nColumns} in all" + ) + raise ValueError(msg) + if len(listed) < widest: + msg = ( + f"Cluster {clusterID} lists {len(listed)} columns, fewer " + f"than the {widest} of cluster {widestClusterID} before it" + ) + raise ValueError(msg) + if len(listed) > widest: + widestClusterID, widest = clusterID, len(listed) + clusters.append(context.cluster(clusterID, summary, listed)) + envelopes.append(clusters) + return envelopes + + +class _ColumnContext: + """What ``get_extended_page_descriptions`` needs from the schema, computed once""" + + def __init__(self, schema: SchemaDescription): + self.schema = schema + # The columns of a field's representations correspond by their index + # within the representation, which ROOT counts in column-ID order + # (tree/ntuple/src/RNTupleSerialize.cxx:1511 at 6.40.04). Each list is in + # column-ID order, so its first column is the first representation's. + self.corresponding: list[list[int]] = [] + groups: dict[tuple[int, int], list[int]] = {} + counts: dict[tuple[int, int], int] = {} + for columnID, column in enumerate(schema.columnDescriptions): + representation = (column.fFieldID, column.fRepresentationIndex) + index = counts.get(representation, 0) + counts[representation] = index + 1 + group = groups.setdefault((column.fFieldID, index), []) + group.append(columnID) + self.corresponding.append(group) + self.paths: dict[int, bytes] = {} + + def path(self, fieldID: int) -> bytes: + if fieldID not in self.paths: + self.paths[fieldID] = self.schema.field_path(fieldID) + return self.paths[fieldID] + + def deferred_range(self, columnID: int, summary: ClusterSummary) -> tuple[int, int]: + """The element range in a cluster of a deferred column of the first representation + + The column holds a fixed number of elements per entry: the product of + the array sizes of its field and its ancestors. The spec allows an + unsuppressed deferred column only where no ancestor is a collection or a + variant (*Column Description*); otherwise the range cannot be computed. + """ + fields = self.schema.fieldDescriptions + fieldID = self.schema.columnDescriptions[columnID].fFieldID + chain = self.schema.field_chain(fieldID) + for fid in chain[1:]: + if fields[fid].fStructuralRole in (0x01, 0x03): # collection, variant + msg = ( + f"Column {columnID} is deferred, but its field " + f"{self.path(fieldID)!r} is inside the collection or variant " + f"{self.path(fid)!r}" + ) + raise ValueError(msg) + repetitions = 1 + for fid in chain: + repetitions *= max(fields[fid].fArraySize or 0, 1) + return ( + summary.fFirstEntryNumber * repetitions, + summary.fNEntries * repetitions, + ) + + def cluster( + self, + clusterID: int, + summary: ClusterSummary, + listed: list[PageLocations[RPageDescription]], + ) -> list[InterpretableColumn]: + """The columns of one cluster, built in the three steps of ROOT's reader""" + columns = self.schema.columnDescriptions + locations: list[PageLocations[RPageDescription] | None] = [ + *listed, + *[None] * (len(columns) - len(listed)), + ] + + # 1. The page list: an unsuppressed column starts at its element offset + # and holds its pages' elements; a column the cluster predates holds none + suppressed: list[bool] = [] + ranges: list[tuple[int, int] | None] = [] + for columnID, (location, column) in enumerate( + zip(locations, columns, strict=True) + ): + if location is None: + suppressed.append((column.fFirstElementIndex or 0) < 0) + ranges.append((0, 0)) + elif location.elementoffset < 0: + if location.items: + msg = ( + f"Cluster {clusterID}: column {columnID} is suppressed " + f"but has {len(location.items)} pages" ) raise ValueError(msg) - # Columns added by a later model extension, which the cluster predates - missing: list[list[RPageDescription]] = [ - [] for _ in range(len(columns) - len(columnlist)) - ] - clusterPages: list[list[InterpretablePage]] = [] - for columnID, (pagelist, column) in enumerate( - zip([*columnlist, *missing], columns, strict=True) + suppressed.append(True) + ranges.append(None) + else: + suppressed.append(False) + nStored = sum(page.n_elements for page in location.items) + ranges.append((location.elementoffset, nStored)) + + # 2. A suppressed column takes the range of the corresponding column of + # the field's active representation (spec, *Suppressed Columns*) + for columnID, column in enumerate(columns): + if ranges[columnID] is not None: + continue + for other in self.corresponding[columnID]: + if ( + columns[other].fRepresentationIndex != column.fRepresentationIndex + and locations[other] is not None + and not suppressed[other] ): - fieldID = column.fFieldID - if fieldID not in paths: - paths[fieldID] = schema.field_path(fieldID) - clusterPages.append( - [ - InterpretablePage( - pageDescription=page_description, - uncompressedSize=ceil( - page_description.n_elements - * column.fBitsOnStorage - / 8 - ), # Convert bits to bytes - columnType=column.fColumnType, - clusterID=clusterID, - columnID=columnID, - fieldID=fieldID, - fieldDescription=schema.fieldDescriptions[fieldID], - fieldPath=paths[fieldID], - ) - for page_description in pagelist - ] + ranges[columnID] = ranges[other] + break + else: + msg = ( + f"Cluster {clusterID}: column {columnID} is suppressed, and no " + f"other representation of field {self.path(column.fFieldID)!r} " + "is active" + ) + raise ValueError(msg) + + # 3. A deferred column covers the whole cluster, the elements before its + # first stored one being zeros; a later representation copies the range + # of the first + for columnID, column in enumerate(columns): + if not column.fFirstElementIndex: + continue + first = self.corresponding[columnID][0] + if column.fRepresentationIndex == 0: + ranges[columnID] = self.deferred_range(columnID, summary) + elif first != columnID: + ranges[columnID] = ranges[first] + else: + msg = ( + f"Column {columnID} is deferred and of representation " + f"{column.fRepresentationIndex}, but field " + f"{self.path(column.fFieldID)!r} has no earlier representation" + ) + raise ValueError(msg) + + out: list[InterpretableColumn] = [] + for columnID, (location, column) in enumerate( + zip(locations, columns, strict=True) + ): + range_ = ranges[columnID] + assert range_ is not None + firstElementIndex, nElements = range_ + descriptions: list[RPageDescription] = [] + nZeroElements = 0 + if location is not None and not suppressed[columnID]: + descriptions = location.items + nStored = sum(page.n_elements for page in descriptions) + nZeroElements = nElements - nStored + if ( + nZeroElements < 0 + or location.elementoffset != firstElementIndex + nZeroElements + ): + msg = ( + f"Cluster {clusterID}: column {columnID} has elements " + f"{firstElementIndex} to {firstElementIndex + nElements}, " + f"but its pages hold {nStored} elements from " + f"{location.elementoffset}" ) - clusters.append(clusterPages) - envelopePages.append(clusters) - return envelopePages + raise ValueError(msg) + elif not suppressed[columnID]: + nZeroElements = nElements + fieldID = column.fFieldID + fieldDescription = self.schema.fieldDescriptions[fieldID] + fieldPath = self.path(fieldID) + out.append( + InterpretableColumn( + clusterID=clusterID, + columnID=columnID, + columnDescription=column, + fieldID=fieldID, + fieldDescription=fieldDescription, + fieldPath=fieldPath, + pageLocations=location, + suppressed=suppressed[columnID], + firstElementIndex=firstElementIndex, + nElements=nElements, + nZeroElements=nZeroElements, + pages=[ + InterpretablePage( + pageDescription=page_description, + uncompressedSize=ceil( + page_description.n_elements * column.fBitsOnStorage / 8 + ), # Convert bits to bytes + columnType=column.fColumnType, + clusterID=clusterID, + columnID=columnID, + fieldID=fieldID, + fieldDescription=fieldDescription, + fieldPath=fieldPath, + ) + for page_description in descriptions + ], + ) + ) + return out diff --git a/tests/test_extended_page_descriptions.py b/tests/test_extended_page_descriptions.py index da4d357..ebb133c 100644 --- a/tests/test_extended_page_descriptions.py +++ b/tests/test_extended_page_descriptions.py @@ -1,14 +1,21 @@ import dataclasses +import re from pathlib import Path import pytest +import tomli # tomllib, once Python 3.10 is dropped from skhep_testdata import data_path # type: ignore[import-not-found] from rootfilespec.reader import open_path from rootfilespec.rntuple.RNTuple import RNTuple, SchemaDescription from rootfilespec.rntuple.schema import ColumnType -DATA = Path(__file__).parent.parent / "reference" / "root-io-spec" / "data" / "rntuple" +SPEC = Path(__file__).parent.parent / "reference" / "root-io-spec" +DATA = SPEC / "data" / "rntuple" + +MULTIPLE_REPRESENTATIONS = "test_multiple_representations_rntuple_v1-0-0-0.root" +MULTIPLE_CLUSTER_GROUPS = "test_multiple_cluster_groups_rntuple_v1-0-0-0.root" +EXTENSION_COLUMNS = "test_extension_columns_rntuple_v1-0-0-0.root" def _load(path: str | Path) -> RNTuple: @@ -18,33 +25,54 @@ def _load(path: str | Path) -> RNTuple: return RNTuple.from_anchor(reader.fetch(keylist[name]), reader.fetch.buffer) +def _ranges(rntuple: RNTuple) -> list[list[tuple[bool, int, int, int, int]]]: + """(suppressed, firstElementIndex, nElements, nZeroElements, pages) per column, per cluster""" + return [ + [ + ( + column.suppressed, + column.firstElementIndex, + column.nElements, + column.nZeroElements, + len(column.pages), + ) + for column in cluster + ] + for envelope in rntuple.get_extended_page_descriptions() + for cluster in envelope + ] + + def test_suppressed_columns_keep_their_position(): """A column's position is its column ID, suppressed or not In this file the field "real" has two representations, Real32 (column 0) and Real16 (column 1), and each cluster suppresses one of them. Before, - cluster 1's only entry was column 1, at position 0. + cluster 1's only entry was column 1, at position 0. A suppressed column + takes the element range of the active one (spec, *Suppressed Columns*). """ - rntuple = _load(data_path("test_multiple_representations_rntuple_v1-0-0-0.root")) - (envelope,) = rntuple.get_extended_page_descriptions() - assert [[len(column) for column in cluster] for cluster in envelope] == [ - [1, 0], - [0, 1], - [1, 0], + rntuple = _load(data_path(MULTIPLE_REPRESENTATIONS)) + assert _ranges(rntuple) == [ + [(False, 0, 1, 0, 1), (True, 0, 1, 0, 0)], + [(True, 1, 1, 0, 0), (False, 1, 1, 0, 1)], + [(False, 2, 1, 0, 1), (True, 2, 1, 0, 0)], ] - (page,) = envelope[1][1] + (envelope,) = rntuple.get_extended_page_descriptions() + (page,) = envelope[1][1].pages assert page.columnID == 1 assert page.columnType == ColumnType.kReal16 suppressed = rntuple.pagelistEnvelopes[0].pageLocations[1][0] assert suppressed.elementoffset < 0 + assert envelope[1][0].pageLocations is suppressed def test_cluster_ids_continue_across_cluster_groups(): """Issue #116: 3 cluster groups of 5, 4 and 3 clusters are clusters 0 to 11""" - rntuple = _load(data_path("test_multiple_cluster_groups_rntuple_v1-0-0-0.root")) + rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) assert rntuple.firstClusterIDs == [0, 5, 9] ids = [ - {page.clusterID for column in cluster for page in column} + {column.clusterID for column in cluster} + | {page.clusterID for column in cluster for page in column.pages} for envelope in rntuple.get_extended_page_descriptions() for cluster in envelope ] @@ -74,9 +102,35 @@ def test_field_paths_user_class(): assert path.rsplit(b".", 1)[-1] == names[i] +@pytest.mark.skipif(not DATA.exists(), reason="reference/root-io-spec not checked out") +def test_columns_user_class(): + """Each column's field path and type, as root-io-spec's case pins them + (": column type = ", at the column's offset)""" + case = tomli.loads( + (SPEC / "gen" / "cases" / "rntuple" / "user-class" / "case.toml").read_text() + ) + pinned = sorted( + (record["offset"], match[1].encode(), ColumnType(record["value"])) + for record in case["bytes"] + if ( + match := re.fullmatch( + r"(\S+): column type \w+ = 0x[0-9a-f]+", record["name"] + ) + ) + ) + assert len(pinned) == 19 + rntuple = _load(DATA / "user-class.root") + for envelope in rntuple.get_extended_page_descriptions(): + for cluster in envelope: + assert [ + (column.fieldPath, column.columnDescription.fColumnType) + for column in cluster + ] == [(path, type_) for _, path, type_ in pinned] + + def test_field_paths_schema_extension(): """Field IDs continue from the header into the footer's schema extension""" - rntuple = _load(data_path("test_extension_columns_rntuple_v1-0-0-0.root")) + rntuple = _load(data_path(EXTENSION_COLUMNS)) assert len(rntuple.headerEnvelope.fieldDescriptions) == 1 schema = rntuple.schemaDescription assert [schema.field_path(i) for i in range(4)] == [ @@ -87,42 +141,58 @@ def test_field_paths_schema_extension(): ] -# scikit-hep-testdata files whose page lists differ from the root-io-spec ones: -# suppressed columns, several cluster groups, and a model extended after the -# first cluster was committed -SKHEP_FILES = [ - "test_multiple_representations_rntuple_v1-0-0-0.root", - "test_multiple_cluster_groups_rntuple_v1-0-0-0.root", - "test_extension_columns_rntuple_v1-0-0-0.root", -] - - def test_columns_added_by_model_extension(): """A cluster committed before the model was extended lists only the columns - that existed then: here 2 of the 4. The two later ones get empty entries, - as ROOT synthesizes them (AddExtendedColumnRanges), so a position is still - the column ID.""" - rntuple = _load(data_path("test_extension_columns_rntuple_v1-0-0-0.root")) + that existed then: here 2 of the 4. The view still has all 4, with the + ranges ROOT's reader gives them (AddExtendedColumnRanges). + + float_field (column 1) and intvec_field (column 2) are deferred from + elements 200 and 400: they cover every cluster from element 0, the elements + before those being zeros with no page (spec, *Column Description*). + intvec_field._0 (column 3) is inside the vector, so it is not deferred: the + zero vectors are empty, and it holds no elements before cluster 1. + """ + rntuple = _load(data_path(EXTENSION_COLUMNS)) + columns = rntuple.schemaDescription.columnDescriptions + assert [c.fFirstElementIndex for c in columns] == [None, 200, 400, None] (pagelist,) = rntuple.pagelistEnvelopes assert [len(columns) for columns in pagelist.pageLocations] == [2, 4, 4, 4] + assert [s.fNEntries for s in pagelist.clusterSummaries] == [350, 117, 84, 49] + assert _ranges(rntuple) == [ + [ + (False, 0, 350, 0, 2), + (False, 0, 350, 200, 1), + (False, 0, 350, 350, 0), + (False, 0, 0, 0, 0), + ], + [ + (False, 350, 117, 0, 1), + (False, 350, 117, 0, 1), + (False, 350, 117, 50, 1), + (False, 0, 134, 0, 1), + ], + [ + (False, 467, 84, 0, 1), + (False, 467, 84, 0, 1), + (False, 467, 84, 0, 1), + (False, 134, 168, 0, 1), + ], + [ + (False, 551, 49, 0, 1), + (False, 551, 49, 0, 1), + (False, 551, 49, 0, 1), + (False, 302, 98, 0, 1), + ], + ] (envelope,) = rntuple.get_extended_page_descriptions() - assert [[len(column) for column in cluster] for cluster in envelope] == [ - [2, 1, 0, 0], - [1, 1, 1, 1], - [1, 1, 1, 1], - [1, 1, 1, 1], + assert [column.pageLocations is None for column in envelope[0]] == [ + False, + False, + True, + True, ] -def test_page_list_with_a_missing_header_column_raises(): - """Only columns of the schema extension may be missing from a page list""" - rntuple = _load(data_path("test_extension_columns_rntuple_v1-0-0-0.root")) - (pagelist,) = rntuple.pagelistEnvelopes - pagelist.pageLocations.items[0].items = [] - with pytest.raises(ValueError, match="Cluster 0 lists 0 columns"): - rntuple.get_extended_page_descriptions() - - @pytest.mark.parametrize( "path", [ @@ -130,10 +200,20 @@ def test_page_list_with_a_missing_header_column_raises(): pytest.param(DATA / name, id=name) for name in sorted(p.name for p in DATA.glob("*.root")) ), - *(pytest.param(name, id=name) for name in SKHEP_FILES), + *( + pytest.param(name, id=name) + for name in [ + MULTIPLE_REPRESENTATIONS, + MULTIPLE_CLUSTER_GROUPS, + EXTENSION_COLUMNS, + ] + ), ], ) -def test_every_page_knows_its_column_and_field(path: Path | str): +def test_every_column_is_complete(path: Path | str): + """In every cluster: one entry per column; a suppressed column has no pages; + the pages hold the elements after the zeros; and a column's clusters follow + each other with no gap, from element 0.""" if isinstance(path, Path): if not path.exists(): pytest.skip("reference/root-io-spec not checked out") @@ -141,22 +221,132 @@ def test_every_page_knows_its_column_and_field(path: Path | str): else: rntuple = _load(data_path(path)) schema = rntuple.schemaDescription + nextElement = [0] * len(schema.columnDescriptions) for envelope in rntuple.get_extended_page_descriptions(): for cluster in envelope: - assert len(cluster) == len(schema.columnDescriptions) - for columnID, column in enumerate(cluster): - for page in column: - assert page.columnID == columnID - field_id = schema.columnDescriptions[columnID].fFieldID - assert page.fieldID == field_id - assert page.fieldDescription == schema.fieldDescriptions[field_id] - assert page.fieldPath == schema.field_path(field_id) + assert [column.columnID for column in cluster] == list( + range(len(schema.columnDescriptions)) + ) + for column in cluster: + if column.suppressed: + assert column.pages == [] + assert column.nZeroElements == 0 + else: + assert column.nElements - column.nZeroElements == sum( + page.pageDescription.n_elements for page in column.pages + ) + assert column.firstElementIndex == nextElement[column.columnID] + nextElement[column.columnID] += column.nElements + for page in column.pages: + assert (page.clusterID, page.columnID, page.fieldPath) == ( + column.clusterID, + column.columnID, + column.fieldPath, + ) + + +def test_page_list_with_a_missing_header_column_raises(): + """Only columns of the schema extension may be missing from a page list""" + rntuple = _load(data_path(EXTENSION_COLUMNS)) + (pagelist,) = rntuple.pagelistEnvelopes + pagelist.pageLocations.items[0].items = [] + with pytest.raises(ValueError, match="Cluster 0 lists 0 columns"): + rntuple.get_extended_page_descriptions() + + +def test_page_list_with_an_extra_column_raises(): + rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) + columnlist = rntuple.pagelistEnvelopes[1].pageLocations.items[2] + columnlist.items.append(columnlist.items[0]) + with pytest.raises(ValueError, match="Cluster 7 lists 4 columns"): + rntuple.get_extended_page_descriptions() + + +def test_page_list_with_fewer_columns_than_an_earlier_one_raises(): + """Only a model extension adds columns, so a later cluster cannot list fewer""" + rntuple = _load(data_path(EXTENSION_COLUMNS)) + (pagelist,) = rntuple.pagelistEnvelopes + pagelist.pageLocations.items[2].items.pop() + with pytest.raises( + ValueError, match="Cluster 2 lists 3 columns, fewer than the 4 of cluster 1" + ): + rntuple.get_extended_page_descriptions() + + +def test_page_lists_must_match_the_cluster_groups(): + rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) + rntuple.pagelistEnvelopes.pop() + with pytest.raises(ValueError, match="2 page lists for 3 cluster groups"): + rntuple.get_extended_page_descriptions() + + +def test_page_list_must_have_its_groups_clusters(): + rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) + pagelist = rntuple.pagelistEnvelopes[1] + pagelist.pageLocations.items.pop() + pagelist.clusterSummaries.items.pop() + with pytest.raises( + ValueError, match="page locations for 3 clusters, the group says 4" + ): + rntuple.get_extended_page_descriptions() + + +def test_page_list_must_have_a_summary_per_cluster(): + rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) + rntuple.pagelistEnvelopes[1].clusterSummaries.items.pop() + with pytest.raises(ValueError, match="has 3 cluster summaries"): + rntuple.get_extended_page_descriptions() + + +def test_suppressed_column_with_pages_raises(): + """Suppressed columns always have an empty list of pages (spec)""" + rntuple = _load(data_path(MULTIPLE_REPRESENTATIONS)) + clusters = rntuple.pagelistEnvelopes[0].pageLocations.items + clusters[1].items[0].items = list(clusters[0].items[0].items) + with pytest.raises(ValueError, match="column 0 is suppressed but has 1 pages"): + rntuple.get_extended_page_descriptions() + + +def test_field_with_every_representation_suppressed_raises(): + """Every field has exactly one active representation in a cluster (spec)""" + rntuple = _load(data_path(MULTIPLE_REPRESENTATIONS)) + clusters = rntuple.pagelistEnvelopes[0].pageLocations.items + clusters[1].items[1] = clusters[1].items[0] + with pytest.raises(ValueError, match="no other representation of field b'real'"): + rntuple.get_extended_page_descriptions() + + +def test_deferred_column_whose_pages_disagree_raises(): + """intvec_field starts at element 400 (column description), so in cluster 1 + (elements 350 to 467) its pages start at 400""" + rntuple = _load(data_path(EXTENSION_COLUMNS)) + rntuple.pagelistEnvelopes[0].pageLocations.items[1].items[2].elementoffset = 401 + with pytest.raises( + ValueError, match="column 2 has elements 350 to 467, but its pages hold 67" + ): + rntuple.get_extended_page_descriptions() + + +def test_deferred_column_inside_a_collection_raises(): + """The format cannot say how many elements a deferred column inside a + collection has (spec, *Column Description*)""" + rntuple = _load(data_path(EXTENSION_COLUMNS)) + extension = rntuple.footerEnvelope.schemaExtension.columnDescriptions.items + extension[2] = dataclasses.replace(extension[2], fFlags=1, fFirstElementIndex=5) + with pytest.raises(ValueError, match="inside the collection or variant"): + rntuple.get_extended_page_descriptions() + + +def test_deferred_column_without_a_first_representation_raises(): + rntuple = _load(data_path(EXTENSION_COLUMNS)) + extension = rntuple.footerEnvelope.schemaExtension.columnDescriptions.items + extension[0] = dataclasses.replace(extension[0], fRepresentationIndex=1) + with pytest.raises(ValueError, match="has no earlier representation"): + rntuple.get_extended_page_descriptions() def test_field_path_errors(): - schema = _load( - data_path("test_extension_columns_rntuple_v1-0-0-0.root") - ).schemaDescription + schema = _load(data_path(EXTENSION_COLUMNS)).schemaDescription with pytest.raises(ValueError, match="parent chain reaches field ID 9"): schema.field_path(9) fields = list(schema.fieldDescriptions) diff --git a/tests/test_rntuple_hardcoded.py b/tests/test_rntuple_hardcoded.py index caac81f..3b0ab00 100644 --- a/tests/test_rntuple_hardcoded.py +++ b/tests/test_rntuple_hardcoded.py @@ -339,7 +339,10 @@ def test_read_contributors(): ) extended_page_descriptions = rntuple.get_extended_page_descriptions() - assert extended_page_descriptions == [ # PagelistEnvelopes + assert [ + [[column.pages for column in cluster] for cluster in envelope] + for envelope in extended_page_descriptions + ] == [ # PagelistEnvelopes [ # Clusters (columnlists) [ # Columns (pagelists) [ # Pages (page descriptions) @@ -615,7 +618,10 @@ def fetch_from_locator(loc): ) extended_page_descriptions_a = rntuple_a.get_extended_page_descriptions() - assert extended_page_descriptions_a == [ # PagelistEnvelopes + assert [ + [[column.pages for column in cluster] for cluster in envelope] + for envelope in extended_page_descriptions_a + ] == [ # PagelistEnvelopes [ # Clusters (columnlists) [ # Columns (pagelists) [ # Pages (page descriptions) @@ -815,7 +821,10 @@ def fetch_from_locator(loc): ) extended_page_descriptions_b = rntuple_b.get_extended_page_descriptions() - assert extended_page_descriptions_b == [ # PagelistEnvelopes + assert [ + [[column.pages for column in cluster] for cluster in envelope] + for envelope in extended_page_descriptions_b + ] == [ # PagelistEnvelopes [ # Clusters (columnlists) [ # Columns (pagelists) [ # Pages (page descriptions) From 1ddc44cbd8d77239077123a27b3068a34e2288a7 Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:55:03 -0500 Subject: [PATCH 6/9] Return clusters, each with its columns, each with its pages RNTuple.clusters() replaces get_extended_page_descriptions(): one InterpretableCluster per cluster, in cluster-ID order, shaped like ROOT's RClusterDescriptor. A cluster carries its ID, its cluster group's position and its summary as stored (first entry, number of entries), so a caller no longer zips the view back with the page lists and the footer. The nesting by page-list envelope goes, and with it RNTuple.firstClusterIDs. Each fact is now stored once. InterpretableColumn keeps the column's and field's descriptions and path; InterpretablePage keeps only the page as stored, its uncompressed size and, new, the column-wide index of its first element (ROOT's RPageInfoExtended). The page's columnType, clusterID, columnID, fieldID, fieldDescription and fieldPath are on its column. SchemaDescription.field_columns() gives every field's physical columns by representation and index, as ROOT numbers them (RNTupleSerialize.cxx:1511 at 6.40.04), and raises when a field's representations differ in size (spec, Suppressed Columns). The cluster view uses it to find corresponding columns. Assisted-by: claude-code:claude-opus-5-5 --- src/rootfilespec/rntuple/RNTuple.py | 337 ++++++++++-------- ...scriptions.py => test_rntuple_clusters.py} | 164 +++++---- tests/test_rntuple_hardcoded.py | 208 +++++------ 3 files changed, 363 insertions(+), 346 deletions(-) rename tests/{test_extended_page_descriptions.py => test_rntuple_clusters.py} (72%) diff --git a/src/rootfilespec/rntuple/RNTuple.py b/src/rootfilespec/rntuple/RNTuple.py index 05dd26d..0a52e45 100644 --- a/src/rootfilespec/rntuple/RNTuple.py +++ b/src/rootfilespec/rntuple/RNTuple.py @@ -15,7 +15,6 @@ from rootfilespec.rntuple.schema import ( AliasColumnDescription, ColumnDescription, - ColumnType, ExtraTypeInformation, FieldDescription, ) @@ -79,7 +78,7 @@ def from_envelopes( extraTypeInformations, ) - def field_chain(self, field_id: int) -> list[int]: + def _field_chain(self, field_id: int) -> list[int]: """The IDs of a field and of its ancestors, from the field up to its top-level field A field ID is the field's position in the combined list (header first, @@ -109,59 +108,86 @@ def field_path(self, field_id: int) -> bytes: "." is forbidden in field names (spec, *Naming specification*), so the result is unambiguous. Names are bytes, as stored. """ - chain = self.field_chain(field_id) + chain = self._field_chain(field_id) return b".".join(self.fieldDescriptions[fid].fFieldName for fid in chain[::-1]) + def field_columns(self) -> list[list[list[int]]]: + """The physical column IDs of every field, by representation and then by index + + ``field_columns()[fieldID][r][i]`` is the ``i``-th column of the field's + representation ``r``, as ROOT numbers them: in column-ID order within + each representation (``RNTupleSerialize.cxx:1511`` at 6.40.04). A + field's representations have the same number of columns, which + correspond one to one (spec, *Suppressed Columns*). A field with no + columns has no representations. + + Raises if a field's representation indices are not 0, 1, ..., or if its + representations have different numbers of columns. + """ + byField: list[dict[int, list[int]]] = [{} for _ in self.fieldDescriptions] + for columnID, column in enumerate(self.columnDescriptions): + if not 0 <= column.fFieldID < len(byField): + msg = f"Column {columnID} belongs to field {column.fFieldID}, of {len(byField)} fields" + raise ValueError(msg) + representations = byField[column.fFieldID] + representations.setdefault(column.fRepresentationIndex, []).append(columnID) + out: list[list[list[int]]] = [] + for fieldID, representations in enumerate(byField): + if sorted(representations) != list(range(len(representations))): + msg = f"Field {fieldID} has representations {sorted(representations)}" + raise ValueError(msg) + ordered = [representations[r] for r in range(len(representations))] + if len({len(columns) for columns in ordered}) > 1: + msg = ( + f"Field {fieldID}: its representations have " + f"{[len(columns) for columns in ordered]} columns" + ) + raise ValueError(msg) + out.append(ordered) + return out + @dataclasses.dataclass class InterpretablePage: - """A class representing an interpretable page description. - It provides the page description, uncompressed size, and column type. + """One page of a column in a cluster + + Like ROOT's ``RClusterDescriptor::RPageInfoExtended``: the page as stored, + and where its elements are in the column. """ pageDescription: RPageDescription - """The RPageDescription object representing the page.""" + """The page's entry in the page list, as stored: its locator, its number of + elements and whether a checksum follows it.""" + firstElementIndex: int + """The index, within the column, of the page's first element.""" uncompressedSize: int - """The uncompressed size of the page, in bytes.""" - columnType: ColumnType - """The type of the column this page belongs to, e.g. kInt32, kFloat64, etc.""" - clusterID: int - """The RNTuple-wide ID of the page's cluster, continuous across cluster groups.""" - columnID: int - """The ID of the page's physical column (its position in the combined column list).""" - fieldID: int - """The ID of the field the column belongs to (``ColumnDescription.fFieldID``).""" - fieldDescription: FieldDescription - """The description of that field.""" - fieldPath: bytes - """The field's qualified name (see ``SchemaDescription.field_path``).""" + """The size of the page's elements packed for storage, before compression, in bytes.""" @dataclasses.dataclass class InterpretableColumn: """One physical column in one cluster: its element range and its pages - Enough on its own to read the column's elements in the cluster. The element - range is the one ROOT's reader builds (``CommitSuppressedColumnRanges`` and - ``AddExtendedColumnRanges``, ``tree/ntuple/src/RNTupleDescriptor.cxx:880`` - and ``:920`` at 6.40.04), which is what the spec asks of a reader - (*Suppressed Columns*; *Column Description* for deferred columns). + Like ROOT's ``RClusterDescriptor::RColumnRange`` with its ``RPageRange``, + and enough on its own to read the column's elements in the cluster. The + element range is the one ROOT's reader builds + (``CommitSuppressedColumnRanges`` and ``AddExtendedColumnRanges``, + ``tree/ntuple/src/RNTupleDescriptor.cxx:880`` and ``:920`` at 6.40.04), + which is what the spec asks of a reader (*Suppressed Columns*; *Column + Description* for deferred columns). """ - clusterID: int - """The RNTuple-wide ID of the cluster, continuous across cluster groups.""" columnID: int """The ID of the physical column (its position in the combined column list).""" columnDescription: ColumnDescription - """The description of the column.""" - fieldID: int - """The ID of the field the column belongs to (``ColumnDescription.fFieldID``).""" + """The description of the column; its field is ``columnDescription.fFieldID``.""" fieldDescription: FieldDescription - """The description of that field.""" + """The description of the column's field.""" fieldPath: bytes """The field's qualified name (see ``SchemaDescription.field_path``).""" pageLocations: PageLocations[RPageDescription] | None - """The column's entry in the cluster's page list, as stored. + """The column's entry in the cluster's page list, as stored: its pages, + element offset and compression settings. ``None`` if the page list has no entry for the column: a cluster committed before the model was extended lists only the columns that existed then @@ -188,6 +214,25 @@ class InterpretableColumn: """The column's pages on disk in this cluster, in order.""" +@dataclasses.dataclass +class InterpretableCluster: + """One cluster: its entries, and every column's elements and pages in it + + Like ROOT's ``RClusterDescriptor``. + """ + + clusterID: int + """The RNTuple-wide ID of the cluster: cluster IDs continue across cluster + groups (spec, *Page List Envelope*).""" + clusterGroupID: int + """The position of the cluster's group in the footer's list of cluster groups.""" + summary: ClusterSummary + """The cluster's summary, as stored: its first entry and number of entries.""" + columns: list[InterpretableColumn] + """One per physical column, in column-ID order, as the spec guarantees + (*Page Locations*): a position is the column ID.""" + + @dataclasses.dataclass class RNTuple: """A class representing an RNTuple.""" @@ -267,27 +312,10 @@ def streamer_infos(self) -> dict[bytes, TStreamerInfo]: infos[info.fName] = info return infos - @property - def firstClusterIDs(self) -> list[int]: - """The ID of the first cluster of each cluster group, in footer order + def clusters(self) -> list[InterpretableCluster]: + """Every cluster, in cluster-ID order, with every column's elements and pages - Cluster IDs continue across cluster groups: a page list's clusters start - at the number of clusters in all earlier groups (spec, *Page List - Envelope*). ``pagelistEnvelopes`` are in the same order as the groups. - """ - first, out = 0, [] - for group in self.footerEnvelope.clusterGroups: - out.append(first) - first += group.fNClusters - return out - - def get_extended_page_descriptions( - self, - ) -> list[list[list[InterpretableColumn]]]: - """Every column of every cluster, by page-list envelope and cluster. - - Each cluster has one entry per physical column, in column-ID order, so a - position is the column ID (spec, *Page Locations*). Each entry is enough + Each cluster has one entry per physical column, so that each is enough on its own to read the column in that cluster; see ``InterpretableColumn``. @@ -301,28 +329,25 @@ def get_extended_page_descriptions( if len(self.pagelistEnvelopes) != len(groups): msg = f"{len(self.pagelistEnvelopes)} page lists for {len(groups)} cluster groups" raise ValueError(msg) - context = _ColumnContext(schema) - envelopes: list[list[list[InterpretableColumn]]] = [] + columns = _Columns(schema) + clusters: list[InterpretableCluster] = [] # The cluster that lists the most columns so far, and how many: only a # model extension adds columns, so no later cluster lists fewer widestClusterID, widest = 0, 0 - for pagelistEnvelope, group, firstClusterID in zip( - self.pagelistEnvelopes, groups, self.firstClusterIDs, strict=True + for clusterGroupID, (pagelistEnvelope, group) in enumerate( + zip(self.pagelistEnvelopes, groups, strict=True) ): pageLocations = pagelistEnvelope.pageLocations.items summaries = pagelistEnvelope.clusterSummaries.items if not len(pageLocations) == len(summaries) == group.fNClusters: msg = ( - f"Page list of cluster group starting at cluster {firstClusterID} " - f"has {len(summaries)} cluster summaries and page locations for " + f"Page list of cluster group {clusterGroupID} has " + f"{len(summaries)} cluster summaries and page locations for " f"{len(pageLocations)} clusters, the group says {group.fNClusters}" ) raise ValueError(msg) - clusters: list[list[InterpretableColumn]] = [] - for index, (columnlist, summary) in enumerate( - zip(pageLocations, summaries, strict=True) - ): - clusterID = firstClusterID + index + for columnlist, summary in zip(pageLocations, summaries, strict=True): + clusterID = len(clusters) listed = columnlist.items if not nHeaderColumns <= len(listed) <= nColumns: msg = ( @@ -339,72 +364,92 @@ def get_extended_page_descriptions( raise ValueError(msg) if len(listed) > widest: widestClusterID, widest = clusterID, len(listed) - clusters.append(context.cluster(clusterID, summary, listed)) - envelopes.append(clusters) - return envelopes + clusters.append( + InterpretableCluster( + clusterID=clusterID, + clusterGroupID=clusterGroupID, + summary=summary, + columns=columns.in_cluster(clusterID, summary, listed), + ) + ) + return clusters -class _ColumnContext: - """What ``get_extended_page_descriptions`` needs from the schema, computed once""" +@dataclasses.dataclass +class _Column: + """What ``RNTuple.clusters`` needs to know of a column, the same in every cluster""" + + description: ColumnDescription + field: FieldDescription + path: bytes + representations: list[list[int]] + """The columns of the column's field, by representation and index""" + index: int + """The column's index within its representation""" + repetitions: int | None + """For a deferred column of the first representation, its elements per entry""" + + +class _Columns: + """The columns of a schema, and how to find their ranges in a cluster""" def __init__(self, schema: SchemaDescription): - self.schema = schema - # The columns of a field's representations correspond by their index - # within the representation, which ROOT counts in column-ID order - # (tree/ntuple/src/RNTupleSerialize.cxx:1511 at 6.40.04). Each list is in - # column-ID order, so its first column is the first representation's. - self.corresponding: list[list[int]] = [] - groups: dict[tuple[int, int], list[int]] = {} - counts: dict[tuple[int, int], int] = {} - for columnID, column in enumerate(schema.columnDescriptions): - representation = (column.fFieldID, column.fRepresentationIndex) - index = counts.get(representation, 0) - counts[representation] = index + 1 - group = groups.setdefault((column.fFieldID, index), []) - group.append(columnID) - self.corresponding.append(group) - self.paths: dict[int, bytes] = {} - - def path(self, fieldID: int) -> bytes: - if fieldID not in self.paths: - self.paths[fieldID] = self.schema.field_path(fieldID) - return self.paths[fieldID] - - def deferred_range(self, columnID: int, summary: ClusterSummary) -> tuple[int, int]: - """The element range in a cluster of a deferred column of the first representation - - The column holds a fixed number of elements per entry: the product of - the array sizes of its field and its ancestors. The spec allows an - unsuppressed deferred column only where no ancestor is a collection or a - variant (*Column Description*); otherwise the range cannot be computed. + fieldColumns = schema.field_columns() + paths: dict[int, bytes] = {} + self.columns: list[_Column] = [] + for columnID, description in enumerate(schema.columnDescriptions): + fieldID = description.fFieldID + if fieldID not in paths: + paths[fieldID] = schema.field_path(fieldID) + representations = fieldColumns[fieldID] + index = representations[description.fRepresentationIndex].index(columnID) + repetitions = None + if description.fFirstElementIndex and description.fRepresentationIndex == 0: + repetitions = self._repetitions(schema, columnID) + self.columns.append( + _Column( + description=description, + field=schema.fieldDescriptions[fieldID], + path=paths[fieldID], + representations=representations, + index=index, + repetitions=repetitions, + ) + ) + + @staticmethod + def _repetitions(schema: SchemaDescription, columnID: int) -> int: + """The elements per entry of a deferred column of the first representation + + The product of the array sizes of its field and its ancestors. The spec + allows an unsuppressed deferred column only where no ancestor is a + collection or a variant (*Column Description*); there, the number of + elements cannot be known. """ - fields = self.schema.fieldDescriptions - fieldID = self.schema.columnDescriptions[columnID].fFieldID - chain = self.schema.field_chain(fieldID) + fields = schema.fieldDescriptions + fieldID = schema.columnDescriptions[columnID].fFieldID + chain = schema._field_chain(fieldID) for fid in chain[1:]: if fields[fid].fStructuralRole in (0x01, 0x03): # collection, variant msg = ( f"Column {columnID} is deferred, but its field " - f"{self.path(fieldID)!r} is inside the collection or variant " - f"{self.path(fid)!r}" + f"{schema.field_path(fieldID)!r} is inside the collection or " + f"variant {schema.field_path(fid)!r}" ) raise ValueError(msg) repetitions = 1 for fid in chain: repetitions *= max(fields[fid].fArraySize or 0, 1) - return ( - summary.fFirstEntryNumber * repetitions, - summary.fNEntries * repetitions, - ) + return repetitions - def cluster( + def in_cluster( self, clusterID: int, summary: ClusterSummary, listed: list[PageLocations[RPageDescription]], ) -> list[InterpretableColumn]: """The columns of one cluster, built in the three steps of ROOT's reader""" - columns = self.schema.columnDescriptions + columns = self.columns locations: list[PageLocations[RPageDescription] | None] = [ *listed, *[None] * (len(columns) - len(listed)), @@ -418,7 +463,7 @@ def cluster( zip(locations, columns, strict=True) ): if location is None: - suppressed.append((column.fFirstElementIndex or 0) < 0) + suppressed.append((column.description.fFirstElementIndex or 0) < 0) ranges.append((0, 0)) elif location.elementoffset < 0: if location.items: @@ -439,40 +484,30 @@ def cluster( for columnID, column in enumerate(columns): if ranges[columnID] is not None: continue - for other in self.corresponding[columnID]: - if ( - columns[other].fRepresentationIndex != column.fRepresentationIndex - and locations[other] is not None - and not suppressed[other] - ): + for representation in column.representations: + other = representation[column.index] + if locations[other] is not None and not suppressed[other]: ranges[columnID] = ranges[other] break else: msg = ( f"Cluster {clusterID}: column {columnID} is suppressed, and no " - f"other representation of field {self.path(column.fFieldID)!r} " - "is active" + f"other representation of field {column.path!r} is active" ) raise ValueError(msg) # 3. A deferred column covers the whole cluster, the elements before its # first stored one being zeros; a later representation copies the range - # of the first + # of the first, once that is known for columnID, column in enumerate(columns): - if not column.fFirstElementIndex: - continue - first = self.corresponding[columnID][0] - if column.fRepresentationIndex == 0: - ranges[columnID] = self.deferred_range(columnID, summary) - elif first != columnID: - ranges[columnID] = ranges[first] - else: - msg = ( - f"Column {columnID} is deferred and of representation " - f"{column.fRepresentationIndex}, but field " - f"{self.path(column.fFieldID)!r} has no earlier representation" + if column.repetitions is not None: + ranges[columnID] = ( + summary.fFirstEntryNumber * column.repetitions, + summary.fNEntries * column.repetitions, ) - raise ValueError(msg) + for columnID, column in enumerate(columns): + if column.repetitions is None and column.description.fFirstElementIndex: + ranges[columnID] = ranges[column.representations[0][column.index]] out: list[InterpretableColumn] = [] for columnID, (location, column) in enumerate( @@ -500,37 +535,33 @@ def cluster( raise ValueError(msg) elif not suppressed[columnID]: nZeroElements = nElements - fieldID = column.fFieldID - fieldDescription = self.schema.fieldDescriptions[fieldID] - fieldPath = self.path(fieldID) + pages: list[InterpretablePage] = [] + nextElement = firstElementIndex + nZeroElements + for description in descriptions: + pages.append( + InterpretablePage( + pageDescription=description, + firstElementIndex=nextElement, + uncompressedSize=ceil( + description.n_elements + * column.description.fBitsOnStorage + / 8 + ), # Convert bits to bytes + ) + ) + nextElement += description.n_elements out.append( InterpretableColumn( - clusterID=clusterID, columnID=columnID, - columnDescription=column, - fieldID=fieldID, - fieldDescription=fieldDescription, - fieldPath=fieldPath, + columnDescription=column.description, + fieldDescription=column.field, + fieldPath=column.path, pageLocations=location, suppressed=suppressed[columnID], firstElementIndex=firstElementIndex, nElements=nElements, nZeroElements=nZeroElements, - pages=[ - InterpretablePage( - pageDescription=page_description, - uncompressedSize=ceil( - page_description.n_elements * column.fBitsOnStorage / 8 - ), # Convert bits to bytes - columnType=column.fColumnType, - clusterID=clusterID, - columnID=columnID, - fieldID=fieldID, - fieldDescription=fieldDescription, - fieldPath=fieldPath, - ) - for page_description in descriptions - ], + pages=pages, ) ) return out diff --git a/tests/test_extended_page_descriptions.py b/tests/test_rntuple_clusters.py similarity index 72% rename from tests/test_extended_page_descriptions.py rename to tests/test_rntuple_clusters.py index ebb133c..79197c1 100644 --- a/tests/test_extended_page_descriptions.py +++ b/tests/test_rntuple_clusters.py @@ -36,10 +36,9 @@ def _ranges(rntuple: RNTuple) -> list[list[tuple[bool, int, int, int, int]]]: column.nZeroElements, len(column.pages), ) - for column in cluster + for column in cluster.columns ] - for envelope in rntuple.get_extended_page_descriptions() - for cluster in envelope + for cluster in rntuple.clusters() ] @@ -57,32 +56,31 @@ def test_suppressed_columns_keep_their_position(): [(True, 1, 1, 0, 0), (False, 1, 1, 0, 1)], [(False, 2, 1, 0, 1), (True, 2, 1, 0, 0)], ] - (envelope,) = rntuple.get_extended_page_descriptions() - (page,) = envelope[1][1].pages - assert page.columnID == 1 - assert page.columnType == ColumnType.kReal16 + cluster = rntuple.clusters()[1] + assert cluster.columns[1].columnDescription.fColumnType == ColumnType.kReal16 + (page,) = cluster.columns[1].pages + assert page.firstElementIndex == 1 suppressed = rntuple.pagelistEnvelopes[0].pageLocations[1][0] assert suppressed.elementoffset < 0 - assert envelope[1][0].pageLocations is suppressed + assert cluster.columns[0].pageLocations is suppressed def test_cluster_ids_continue_across_cluster_groups(): """Issue #116: 3 cluster groups of 5, 4 and 3 clusters are clusters 0 to 11""" - rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) - assert rntuple.firstClusterIDs == [0, 5, 9] - ids = [ - {column.clusterID for column in cluster} - | {page.clusterID for column in cluster for page in column.pages} - for envelope in rntuple.get_extended_page_descriptions() - for cluster in envelope + clusters = _load(data_path(MULTIPLE_CLUSTER_GROUPS)).clusters() + assert [cluster.clusterID for cluster in clusters] == list(range(12)) + assert [cluster.clusterGroupID for cluster in clusters] == [0] * 5 + [1] * 4 + [ + 2 + ] * 3 + assert [cluster.summary.fFirstEntryNumber for cluster in clusters] == [ + *(0, 100, 200, 300, 400), + *(450, 500, 600, 700), + *(750, 800, 900), ] - assert ids == [{i} for i in range(12)] - first_entries = [ - summary.fFirstEntryNumber - for pagelist in rntuple.pagelistEnvelopes - for summary in pagelist.clusterSummaries + # Column 2 has two elements per entry + assert [cluster.columns[2].pages[0].firstElementIndex for cluster in clusters] == [ + 2 * cluster.summary.fFirstEntryNumber for cluster in clusters ] - assert first_entries == [0, 100, 200, 300, 400, 450, 500, 600, 700, 750, 800, 900] @pytest.mark.skipif(not DATA.exists(), reason="reference/root-io-spec not checked out") @@ -119,13 +117,11 @@ def test_columns_user_class(): ) ) assert len(pinned) == 19 - rntuple = _load(DATA / "user-class.root") - for envelope in rntuple.get_extended_page_descriptions(): - for cluster in envelope: - assert [ - (column.fieldPath, column.columnDescription.fColumnType) - for column in cluster - ] == [(path, type_) for _, path, type_ in pinned] + for cluster in _load(DATA / "user-class.root").clusters(): + assert [ + (column.fieldPath, column.columnDescription.fColumnType) + for column in cluster.columns + ] == [(path, type_) for _, path, type_ in pinned] def test_field_paths_schema_extension(): @@ -141,6 +137,43 @@ def test_field_paths_schema_extension(): ] +@pytest.mark.skipif(not DATA.exists(), reason="reference/root-io-spec not checked out") +def test_field_columns(): + """A field's columns by representation, then by index: a std::string has an + index and a character column; a record has none""" + schema = _load(DATA / "user-class.root").schemaDescription + columns = schema.field_columns() + assert len(columns) == len(schema.fieldDescriptions) + paths = [schema.field_path(i) for i in range(len(columns))] + assert columns[paths.index(b"fHit")] == [] + label = columns[paths.index(b"fHit.fLabel")] + assert [ + [schema.columnDescriptions[c].fColumnType for c in representation] + for representation in label + ] == [[ColumnType.kIndex64, ColumnType.kChar]] + two = _load(data_path(MULTIPLE_REPRESENTATIONS)).schemaDescription + assert two.field_columns() == [[[0], [1]]] + + +def test_field_columns_with_unequal_representations_raise(): + """A field's representations have the same number of columns (spec, + *Suppressed Columns*)""" + schema = _load(data_path(MULTIPLE_REPRESENTATIONS)).schemaDescription + columns = list(schema.columnDescriptions) + columns.append(dataclasses.replace(columns[1])) + broken = dataclasses.replace(schema, columnDescriptions=columns) + with pytest.raises(ValueError, match=r"representations have \[1, 2\] columns"): + broken.field_columns() + + +def test_field_columns_with_a_missing_representation_raise(): + rntuple = _load(data_path(EXTENSION_COLUMNS)) + extension = rntuple.footerEnvelope.schemaExtension.columnDescriptions.items + extension[0] = dataclasses.replace(extension[0], fRepresentationIndex=1) + with pytest.raises(ValueError, match=r"Field 1 has representations \[1\]"): + rntuple.clusters() + + def test_columns_added_by_model_extension(): """A cluster committed before the model was extended lists only the columns that existed then: here 2 of the 4. The view still has all 4, with the @@ -184,13 +217,15 @@ def test_columns_added_by_model_extension(): (False, 302, 98, 0, 1), ], ] - (envelope,) = rntuple.get_extended_page_descriptions() - assert [column.pageLocations is None for column in envelope[0]] == [ + first = rntuple.clusters()[0] + assert [column.pageLocations is None for column in first.columns] == [ False, False, True, True, ] + # float_field's one page in cluster 0 starts after its 200 zeros + assert [page.firstElementIndex for page in first.columns[1].pages] == [200] @pytest.mark.parametrize( @@ -212,37 +247,30 @@ def test_columns_added_by_model_extension(): ) def test_every_column_is_complete(path: Path | str): """In every cluster: one entry per column; a suppressed column has no pages; - the pages hold the elements after the zeros; and a column's clusters follow - each other with no gap, from element 0.""" + the pages follow the zeros and each other with no gap; and a column's + clusters follow each other with no gap, from element 0.""" if isinstance(path, Path): if not path.exists(): pytest.skip("reference/root-io-spec not checked out") rntuple = _load(path) else: rntuple = _load(data_path(path)) - schema = rntuple.schemaDescription - nextElement = [0] * len(schema.columnDescriptions) - for envelope in rntuple.get_extended_page_descriptions(): - for cluster in envelope: - assert [column.columnID for column in cluster] == list( - range(len(schema.columnDescriptions)) - ) - for column in cluster: - if column.suppressed: - assert column.pages == [] - assert column.nZeroElements == 0 - else: - assert column.nElements - column.nZeroElements == sum( - page.pageDescription.n_elements for page in column.pages - ) - assert column.firstElementIndex == nextElement[column.columnID] - nextElement[column.columnID] += column.nElements + nColumns = len(rntuple.schemaDescription.columnDescriptions) + nextElement = [0] * nColumns + for cluster in rntuple.clusters(): + assert [column.columnID for column in cluster.columns] == list(range(nColumns)) + for column in cluster.columns: + if column.suppressed: + assert column.pages == [] + assert column.nZeroElements == 0 + else: + start = column.firstElementIndex + column.nZeroElements for page in column.pages: - assert (page.clusterID, page.columnID, page.fieldPath) == ( - column.clusterID, - column.columnID, - column.fieldPath, - ) + assert page.firstElementIndex == start + start += page.pageDescription.n_elements + assert start == column.firstElementIndex + column.nElements + assert column.firstElementIndex == nextElement[column.columnID] + nextElement[column.columnID] += column.nElements def test_page_list_with_a_missing_header_column_raises(): @@ -251,7 +279,7 @@ def test_page_list_with_a_missing_header_column_raises(): (pagelist,) = rntuple.pagelistEnvelopes pagelist.pageLocations.items[0].items = [] with pytest.raises(ValueError, match="Cluster 0 lists 0 columns"): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_page_list_with_an_extra_column_raises(): @@ -259,7 +287,7 @@ def test_page_list_with_an_extra_column_raises(): columnlist = rntuple.pagelistEnvelopes[1].pageLocations.items[2] columnlist.items.append(columnlist.items[0]) with pytest.raises(ValueError, match="Cluster 7 lists 4 columns"): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_page_list_with_fewer_columns_than_an_earlier_one_raises(): @@ -270,14 +298,14 @@ def test_page_list_with_fewer_columns_than_an_earlier_one_raises(): with pytest.raises( ValueError, match="Cluster 2 lists 3 columns, fewer than the 4 of cluster 1" ): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_page_lists_must_match_the_cluster_groups(): rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) rntuple.pagelistEnvelopes.pop() with pytest.raises(ValueError, match="2 page lists for 3 cluster groups"): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_page_list_must_have_its_groups_clusters(): @@ -288,14 +316,14 @@ def test_page_list_must_have_its_groups_clusters(): with pytest.raises( ValueError, match="page locations for 3 clusters, the group says 4" ): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_page_list_must_have_a_summary_per_cluster(): rntuple = _load(data_path(MULTIPLE_CLUSTER_GROUPS)) rntuple.pagelistEnvelopes[1].clusterSummaries.items.pop() with pytest.raises(ValueError, match="has 3 cluster summaries"): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_suppressed_column_with_pages_raises(): @@ -304,7 +332,7 @@ def test_suppressed_column_with_pages_raises(): clusters = rntuple.pagelistEnvelopes[0].pageLocations.items clusters[1].items[0].items = list(clusters[0].items[0].items) with pytest.raises(ValueError, match="column 0 is suppressed but has 1 pages"): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_field_with_every_representation_suppressed_raises(): @@ -313,7 +341,7 @@ def test_field_with_every_representation_suppressed_raises(): clusters = rntuple.pagelistEnvelopes[0].pageLocations.items clusters[1].items[1] = clusters[1].items[0] with pytest.raises(ValueError, match="no other representation of field b'real'"): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_deferred_column_whose_pages_disagree_raises(): @@ -324,7 +352,7 @@ def test_deferred_column_whose_pages_disagree_raises(): with pytest.raises( ValueError, match="column 2 has elements 350 to 467, but its pages hold 67" ): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_deferred_column_inside_a_collection_raises(): @@ -334,15 +362,7 @@ def test_deferred_column_inside_a_collection_raises(): extension = rntuple.footerEnvelope.schemaExtension.columnDescriptions.items extension[2] = dataclasses.replace(extension[2], fFlags=1, fFirstElementIndex=5) with pytest.raises(ValueError, match="inside the collection or variant"): - rntuple.get_extended_page_descriptions() - - -def test_deferred_column_without_a_first_representation_raises(): - rntuple = _load(data_path(EXTENSION_COLUMNS)) - extension = rntuple.footerEnvelope.schemaExtension.columnDescriptions.items - extension[0] = dataclasses.replace(extension[0], fRepresentationIndex=1) - with pytest.raises(ValueError, match="has no earlier representation"): - rntuple.get_extended_page_descriptions() + rntuple.clusters() def test_field_path_errors(): diff --git a/tests/test_rntuple_hardcoded.py b/tests/test_rntuple_hardcoded.py index 3b0ab00..89819b2 100644 --- a/tests/test_rntuple_hardcoded.py +++ b/tests/test_rntuple_hardcoded.py @@ -338,84 +338,62 @@ def test_read_contributors(): extraTypeInformations=[], ) - extended_page_descriptions = rntuple.get_extended_page_descriptions() + clusters = rntuple.clusters() assert [ - [[column.pages for column in cluster] for cluster in envelope] - for envelope in extended_page_descriptions - ] == [ # PagelistEnvelopes - [ # Clusters (columnlists) - [ # Columns (pagelists) - [ # Pages (page descriptions) - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-22, - locator=StandardLocator(size=176, offset=620), - ), - uncompressedSize=176, - columnType=ColumnType.kIndex64, - clusterID=0, - columnID=0, - fieldID=0, - fieldDescription=rntuple.schemaDescription.fieldDescriptions[ - 0 - ], - fieldPath=b"firstName", - ) - ], - [ - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-178, - locator=StandardLocator(size=178, offset=804), - ), - uncompressedSize=178, - columnType=ColumnType.kChar, - clusterID=0, - columnID=1, - fieldID=0, - fieldDescription=rntuple.schemaDescription.fieldDescriptions[ - 0 - ], - fieldPath=b"firstName", - ) - ], - [ - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-22, - locator=StandardLocator(size=176, offset=990), - ), - uncompressedSize=176, - columnType=ColumnType.kIndex64, - clusterID=0, - columnID=2, - fieldID=1, - fieldDescription=rntuple.schemaDescription.fieldDescriptions[ - 1 - ], - fieldPath=b"lastName", - ) - ], - [ - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-193, - locator=StandardLocator(size=193, offset=1174), - ), - uncompressedSize=193, - columnType=ColumnType.kChar, - clusterID=0, - columnID=3, - fieldID=1, - fieldDescription=rntuple.schemaDescription.fieldDescriptions[ - 1 - ], - fieldPath=b"lastName", - ) - ], - ] + [column.pages for column in cluster.columns] for cluster in clusters + ] == [ # Clusters (columnlists) + [ # Columns (pagelists) + [ # Pages (page descriptions) + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-22, + locator=StandardLocator(size=176, offset=620), + ), + firstElementIndex=0, + uncompressedSize=176, + ) + ], + [ + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-178, + locator=StandardLocator(size=178, offset=804), + ), + firstElementIndex=0, + uncompressedSize=178, + ) + ], + [ + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-22, + locator=StandardLocator(size=176, offset=990), + ), + firstElementIndex=0, + uncompressedSize=176, + ) + ], + [ + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-193, + locator=StandardLocator(size=193, offset=1174), + ), + firstElementIndex=0, + uncompressedSize=193, + ) + ], ] ] + assert [ + (column.columnDescription.fColumnType, column.fieldPath) + for column in clusters[0].columns + ] == [ + (ColumnType.kIndex64, b"firstName"), + (ColumnType.kChar, b"firstName"), + (ColumnType.kIndex64, b"lastName"), + (ColumnType.kChar, b"lastName"), + ] def test_read_multiple_rntuples(): @@ -617,33 +595,27 @@ def fetch_from_locator(loc): extraTypeInformations=[], ) - extended_page_descriptions_a = rntuple_a.get_extended_page_descriptions() + clusters_a = rntuple_a.clusters() assert [ - [[column.pages for column in cluster] for cluster in envelope] - for envelope in extended_page_descriptions_a - ] == [ # PagelistEnvelopes - [ # Clusters (columnlists) - [ # Columns (pagelists) - [ # Pages (page descriptions) - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-100, - locator=StandardLocator(size=138, offset=409), - ), - uncompressedSize=400, - columnType=ColumnType.kSplitReal32, - clusterID=0, - columnID=0, - fieldID=0, - fieldDescription=rntuple_a.schemaDescription.fieldDescriptions[ - 0 - ], - fieldPath=b"f", - ) - ] + [column.pages for column in cluster.columns] for cluster in clusters_a + ] == [ # Clusters (columnlists) + [ # Columns (pagelists) + [ # Pages (page descriptions) + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-100, + locator=StandardLocator(size=138, offset=409), + ), + firstElementIndex=0, + uncompressedSize=400, + ) ] ] ] + assert [ + (column.columnDescription.fColumnType, column.fieldPath) + for column in clusters_a[0].columns + ] == [(ColumnType.kSplitReal32, b"f")] anchor_b = keylist[b"B"].read_object(fetch_data, ROOT3a3aRNTuple) assert anchor_b == ROOT3a3aRNTuple( @@ -820,30 +792,24 @@ def fetch_from_locator(loc): extraTypeInformations=[], ) - extended_page_descriptions_b = rntuple_b.get_extended_page_descriptions() + clusters_b = rntuple_b.clusters() assert [ - [[column.pages for column in cluster] for cluster in envelope] - for envelope in extended_page_descriptions_b - ] == [ # PagelistEnvelopes - [ # Clusters (columnlists) - [ # Columns (pagelists) - [ # Pages (page descriptions) - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-100, - locator=StandardLocator(size=164, offset=1695), - ), - uncompressedSize=400, - columnType=ColumnType.kSplitInt32, - clusterID=0, - columnID=0, - fieldID=0, - fieldDescription=rntuple_b.schemaDescription.fieldDescriptions[ - 0 - ], - fieldPath=b"g", - ) - ] + [column.pages for column in cluster.columns] for cluster in clusters_b + ] == [ # Clusters (columnlists) + [ # Columns (pagelists) + [ # Pages (page descriptions) + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-100, + locator=StandardLocator(size=164, offset=1695), + ), + firstElementIndex=0, + uncompressedSize=400, + ) ] ] ] + assert [ + (column.columnDescription.fColumnType, column.fieldPath) + for column in clusters_b[0].columns + ] == [(ColumnType.kSplitInt32, b"g")] From 0be98ab13fcc937dbfbd0bfd7d3079171975691a Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Fri, 2 Oct 2026 15:06:28 -0500 Subject: [PATCH 7/9] Count a page's first element from the start of its cluster InterpretablePage.firstElementIndex, counted from the start of the column, becomes firstElementInCluster, counted from the start of the column's elements in the cluster, as ROOT's RPageInfoExtended counts it. A cluster can be reused unchanged in another RNTuple, because offset columns count from the start of the cluster (spec, Column Description), so this index belongs to the cluster, while the column's firstElementIndex is a position in this RNTuple. InterpretableCluster's docstring now says which values are which. Assisted-by: claude-code:claude-opus-5-5 --- src/rootfilespec/rntuple/RNTuple.py | 21 +++++++++++++++------ tests/test_rntuple_clusters.py | 25 +++++++++++++++++-------- tests/test_rntuple_hardcoded.py | 12 ++++++------ 3 files changed, 38 insertions(+), 20 deletions(-) diff --git a/src/rootfilespec/rntuple/RNTuple.py b/src/rootfilespec/rntuple/RNTuple.py index 0a52e45..8d4d644 100644 --- a/src/rootfilespec/rntuple/RNTuple.py +++ b/src/rootfilespec/rntuple/RNTuple.py @@ -152,14 +152,18 @@ class InterpretablePage: """One page of a column in a cluster Like ROOT's ``RClusterDescriptor::RPageInfoExtended``: the page as stored, - and where its elements are in the column. + and where its elements start in the cluster. """ pageDescription: RPageDescription """The page's entry in the page list, as stored: its locator, its number of elements and whether a checksum follows it.""" - firstElementIndex: int - """The index, within the column, of the page's first element.""" + firstElementInCluster: int + """The index of the page's first element among the column's elements in the + cluster, as ROOT counts it. + + It belongs to the cluster (see ``InterpretableCluster``); the index within + the column is ``InterpretableColumn.firstElementIndex`` plus this.""" uncompressedSize: int """The size of the page's elements packed for storage, before compression, in bytes.""" @@ -218,7 +222,12 @@ class InterpretableColumn: class InterpretableCluster: """One cluster: its entries, and every column's elements and pages in it - Like ROOT's ``RClusterDescriptor``. + Like ROOT's ``RClusterDescriptor``. ``clusterID``, + ``summary.fFirstEntryNumber`` and each column's ``firstElementIndex`` are + positions in this RNTuple. Everything else belongs to the cluster, which can + be reused unchanged in another RNTuple: offset columns count from the start + of the cluster (spec, *Column Description* and *Stdlib Types and + Collections*). """ clusterID: int @@ -536,12 +545,12 @@ def in_cluster( elif not suppressed[columnID]: nZeroElements = nElements pages: list[InterpretablePage] = [] - nextElement = firstElementIndex + nZeroElements + nextElement = nZeroElements for description in descriptions: pages.append( InterpretablePage( pageDescription=description, - firstElementIndex=nextElement, + firstElementInCluster=nextElement, uncompressedSize=ceil( description.n_elements * column.description.fBitsOnStorage diff --git a/tests/test_rntuple_clusters.py b/tests/test_rntuple_clusters.py index 79197c1..2a6bea4 100644 --- a/tests/test_rntuple_clusters.py +++ b/tests/test_rntuple_clusters.py @@ -59,7 +59,7 @@ def test_suppressed_columns_keep_their_position(): cluster = rntuple.clusters()[1] assert cluster.columns[1].columnDescription.fColumnType == ColumnType.kReal16 (page,) = cluster.columns[1].pages - assert page.firstElementIndex == 1 + assert page.firstElementInCluster == 0 suppressed = rntuple.pagelistEnvelopes[0].pageLocations[1][0] assert suppressed.elementoffset < 0 assert cluster.columns[0].pageLocations is suppressed @@ -77,10 +77,13 @@ def test_cluster_ids_continue_across_cluster_groups(): *(450, 500, 600, 700), *(750, 800, 900), ] - # Column 2 has two elements per entry - assert [cluster.columns[2].pages[0].firstElementIndex for cluster in clusters] == [ + # Column 2 has two elements per entry; its pages count from its cluster's start + assert [cluster.columns[2].firstElementIndex for cluster in clusters] == [ 2 * cluster.summary.fFirstEntryNumber for cluster in clusters ] + assert { + cluster.columns[2].pages[0].firstElementInCluster for cluster in clusters + } == {0} @pytest.mark.skipif(not DATA.exists(), reason="reference/root-io-spec not checked out") @@ -224,8 +227,14 @@ def test_columns_added_by_model_extension(): True, True, ] - # float_field's one page in cluster 0 starts after its 200 zeros - assert [page.firstElementIndex for page in first.columns[1].pages] == [200] + # A first stored page starts after the zeros: float_field's at element 200 of + # cluster 0, and intvec_field's at element 50 of cluster 1, element 400 of + # the column + assert [page.firstElementInCluster for page in first.columns[1].pages] == [200] + column = rntuple.clusters()[1].columns[2] + (page,) = column.pages + assert page.firstElementInCluster == 50 + assert column.firstElementIndex + page.firstElementInCluster == 400 @pytest.mark.parametrize( @@ -264,11 +273,11 @@ def test_every_column_is_complete(path: Path | str): assert column.pages == [] assert column.nZeroElements == 0 else: - start = column.firstElementIndex + column.nZeroElements + start = column.nZeroElements for page in column.pages: - assert page.firstElementIndex == start + assert page.firstElementInCluster == start start += page.pageDescription.n_elements - assert start == column.firstElementIndex + column.nElements + assert start == column.nElements assert column.firstElementIndex == nextElement[column.columnID] nextElement[column.columnID] += column.nElements diff --git a/tests/test_rntuple_hardcoded.py b/tests/test_rntuple_hardcoded.py index 89819b2..b02d4ab 100644 --- a/tests/test_rntuple_hardcoded.py +++ b/tests/test_rntuple_hardcoded.py @@ -349,7 +349,7 @@ def test_read_contributors(): fNElements=-22, locator=StandardLocator(size=176, offset=620), ), - firstElementIndex=0, + firstElementInCluster=0, uncompressedSize=176, ) ], @@ -359,7 +359,7 @@ def test_read_contributors(): fNElements=-178, locator=StandardLocator(size=178, offset=804), ), - firstElementIndex=0, + firstElementInCluster=0, uncompressedSize=178, ) ], @@ -369,7 +369,7 @@ def test_read_contributors(): fNElements=-22, locator=StandardLocator(size=176, offset=990), ), - firstElementIndex=0, + firstElementInCluster=0, uncompressedSize=176, ) ], @@ -379,7 +379,7 @@ def test_read_contributors(): fNElements=-193, locator=StandardLocator(size=193, offset=1174), ), - firstElementIndex=0, + firstElementInCluster=0, uncompressedSize=193, ) ], @@ -606,7 +606,7 @@ def fetch_from_locator(loc): fNElements=-100, locator=StandardLocator(size=138, offset=409), ), - firstElementIndex=0, + firstElementInCluster=0, uncompressedSize=400, ) ] @@ -803,7 +803,7 @@ def fetch_from_locator(loc): fNElements=-100, locator=StandardLocator(size=164, offset=1695), ), - firstElementIndex=0, + firstElementInCluster=0, uncompressedSize=400, ) ] From 3d60c624a01e654bd2ea3c817ac7e6ca040335c7 Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Fri, 2 Oct 2026 15:10:56 -0500 Subject: [PATCH 8/9] Test a column that names a field the schema does not have Assisted-by: claude-code:claude-opus-5-5 --- tests/test_rntuple_clusters.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_rntuple_clusters.py b/tests/test_rntuple_clusters.py index 2a6bea4..8976547 100644 --- a/tests/test_rntuple_clusters.py +++ b/tests/test_rntuple_clusters.py @@ -169,6 +169,15 @@ def test_field_columns_with_unequal_representations_raise(): broken.field_columns() +def test_field_columns_with_a_column_of_no_field_raise(): + schema = _load(data_path(MULTIPLE_REPRESENTATIONS)).schemaDescription + columns = list(schema.columnDescriptions) + columns[1] = dataclasses.replace(columns[1], fFieldID=5) + broken = dataclasses.replace(schema, columnDescriptions=columns) + with pytest.raises(ValueError, match="Column 1 belongs to field 5, of 1 fields"): + broken.field_columns() + + def test_field_columns_with_a_missing_representation_raise(): rntuple = _load(data_path(EXTENSION_COLUMNS)) extension = rntuple.footerEnvelope.schemaExtension.columnDescriptions.items From 9faf8a63fdd5b586a539e1327f957ff441c2f496 Mon Sep 17 00:00:00 2001 From: Samantha Abbott <52800387+sathabbott@users.noreply.github.com> Date: Mon, 5 Oct 2026 12:21:45 -0500 Subject: [PATCH 9/9] Yield clusters one at a time, and make the columns a step of their own Following Nick's review of #131: - RNTuple.clusters() is now an iterator: each cluster is built when it is reached, so a caller that goes through them holds one cluster's ranges at a time. On the MiniAOD Events RNTuple (10,336 columns, 27 clusters), peak memory while iterating is 11 MB, against 84 MB for the whole list. The schema and the number of page lists and clusters are still checked when it is called; each cluster's page list when it is reached. - RNTuple.columns() is the public form of the per-column work that was private (_Columns): one InterpretableColumn per physical column, the same in every cluster, like ROOT's RColumnDescriptor, with its description, field, path and index within its representation. A caller can store these before going through the clusters. - What was InterpretableColumn, one column in one cluster, is now InterpretableColumnRange (ROOT's RColumnRange with its RPageRange). It refers to its InterpretableColumn instead of repeating the column's description, field and path in every cluster, and InterpretableCluster.columns becomes columnRanges. Assisted-by: claude-code:claude-opus-5-5 --- src/rootfilespec/rntuple/RNTuple.py | 317 ++++++++++++++++------------ tests/test_rntuple_clusters.py | 126 +++++++---- tests/test_rntuple_hardcoded.py | 21 +- 3 files changed, 277 insertions(+), 187 deletions(-) diff --git a/src/rootfilespec/rntuple/RNTuple.py b/src/rootfilespec/rntuple/RNTuple.py index 8d4d644..1b62e2c 100644 --- a/src/rootfilespec/rntuple/RNTuple.py +++ b/src/rootfilespec/rntuple/RNTuple.py @@ -1,5 +1,5 @@ import dataclasses -from collections.abc import Callable +from collections.abc import Callable, Iterator from math import ceil from rootfilespec.bootstrap import BOOTSTRAP_CONTEXT @@ -147,6 +147,31 @@ def field_columns(self) -> list[list[list[int]]]: return out +@dataclasses.dataclass +class InterpretableColumn: + """One physical column of the RNTuple, the same in every cluster + + Like ROOT's ``RColumnDescriptor``: the column's description, its field, and + its place among the field's columns. ``RNTuple.columns()`` lists them, and + each cluster's ``InterpretableColumnRange`` refers to one. + """ + + columnID: int + """The ID of the physical column (its position in the combined column list).""" + columnDescription: ColumnDescription + """The description of the column, as stored. Its field is ``fFieldID`` and + its representation ``fRepresentationIndex``.""" + fieldDescription: FieldDescription + """The description of the column's field.""" + fieldPath: bytes + """The field's qualified name (see ``SchemaDescription.field_path``).""" + index: int + """The column's index among the columns of its representation, as ROOT + numbers them (``RColumnDescriptor::GetIndex``). Columns of the field's other + representations with the same index correspond to this one (see + ``SchemaDescription.field_columns``).""" + + @dataclasses.dataclass class InterpretablePage: """One page of a column in a cluster @@ -163,13 +188,13 @@ class InterpretablePage: cluster, as ROOT counts it. It belongs to the cluster (see ``InterpretableCluster``); the index within - the column is ``InterpretableColumn.firstElementIndex`` plus this.""" + the column is ``InterpretableColumnRange.firstElementIndex`` plus this.""" uncompressedSize: int """The size of the page's elements packed for storage, before compression, in bytes.""" @dataclasses.dataclass -class InterpretableColumn: +class InterpretableColumnRange: """One physical column in one cluster: its element range and its pages Like ROOT's ``RClusterDescriptor::RColumnRange`` with its ``RPageRange``, @@ -181,14 +206,8 @@ class InterpretableColumn: Description* for deferred columns). """ - columnID: int - """The ID of the physical column (its position in the combined column list).""" - columnDescription: ColumnDescription - """The description of the column; its field is ``columnDescription.fFieldID``.""" - fieldDescription: FieldDescription - """The description of the column's field.""" - fieldPath: bytes - """The field's qualified name (see ``SchemaDescription.field_path``).""" + column: InterpretableColumn + """The column, as ``RNTuple.columns()`` lists it.""" pageLocations: PageLocations[RPageDescription] | None """The column's entry in the cluster's page list, as stored: its pages, element offset and compression settings. @@ -223,10 +242,10 @@ class InterpretableCluster: """One cluster: its entries, and every column's elements and pages in it Like ROOT's ``RClusterDescriptor``. ``clusterID``, - ``summary.fFirstEntryNumber`` and each column's ``firstElementIndex`` are - positions in this RNTuple. Everything else belongs to the cluster, which can - be reused unchanged in another RNTuple: offset columns count from the start - of the cluster (spec, *Column Description* and *Stdlib Types and + ``summary.fFirstEntryNumber`` and each column range's ``firstElementIndex`` + are positions in this RNTuple. Everything else belongs to the cluster, which + can be reused unchanged in another RNTuple: offset columns count from the + start of the cluster (spec, *Column Description* and *Stdlib Types and Collections*). """ @@ -237,7 +256,7 @@ class InterpretableCluster: """The position of the cluster's group in the footer's list of cluster groups.""" summary: ClusterSummary """The cluster's summary, as stored: its first entry and number of entries.""" - columns: list[InterpretableColumn] + columnRanges: list[InterpretableColumnRange] """One per physical column, in column-ID order, as the spec guarantees (*Page Locations*): a position is the column ID.""" @@ -321,47 +340,136 @@ def streamer_infos(self) -> dict[bytes, TStreamerInfo]: infos[info.fName] = info return infos - def clusters(self) -> list[InterpretableCluster]: + def columns(self) -> list[InterpretableColumn]: + """Every physical column, in column-ID order, with its field and its place among the field's columns + + They are the same in every cluster, so a caller can, say, store them + before going through the clusters with ``clusters()``. Raises + ``ValueError`` where the schema contradicts itself (see + ``SchemaDescription.field_columns``). + """ + schema = self.schemaDescription + return _columns(schema, schema.field_columns()) + + def clusters(self) -> Iterator[InterpretableCluster]: """Every cluster, in cluster-ID order, with every column's elements and pages - Each cluster has one entry per physical column, so that each is enough - on its own to read the column in that cluster; see - ``InterpretableColumn``. + Each cluster has one range per physical column, enough on its own to read + the column in that cluster; see ``InterpretableColumnRange``. Clusters + are built one at a time, as the iterator reaches them, so a caller that + goes through them in turn holds one cluster's ranges at a time. Raises ``ValueError`` where the page lists contradict the footer or the - schema. + schema. The schema, and the number of page lists and of their clusters, + are checked when this is called; each cluster's page list, when the + iterator reaches it. """ - schema = self.schemaDescription - nColumns = len(schema.columnDescriptions) - nHeaderColumns = len(self.headerEnvelope.columnDescriptions.items) - groups = self.footerEnvelope.clusterGroups - if len(self.pagelistEnvelopes) != len(groups): - msg = f"{len(self.pagelistEnvelopes)} page lists for {len(groups)} cluster groups" + return _ClusterReader(self).clusters() + + +def _columns( + schema: SchemaDescription, fieldColumns: list[list[list[int]]] +) -> list[InterpretableColumn]: + """The physical columns of a schema, with their fields and indices""" + paths: dict[int, bytes] = {} + columns: list[InterpretableColumn] = [] + for columnID, description in enumerate(schema.columnDescriptions): + fieldID = description.fFieldID + if fieldID not in paths: + paths[fieldID] = schema.field_path(fieldID) + representation = fieldColumns[fieldID][description.fRepresentationIndex] + columns.append( + InterpretableColumn( + columnID=columnID, + columnDescription=description, + fieldDescription=schema.fieldDescriptions[fieldID], + fieldPath=paths[fieldID], + index=representation.index(columnID), + ) + ) + return columns + + +def _repetitions(schema: SchemaDescription, columnID: int) -> int: + """The elements per entry of a deferred column of the first representation + + The product of the array sizes of its field and its ancestors. The spec + allows an unsuppressed deferred column only where no ancestor is a + collection or a variant (*Column Description*); there, the number of + elements cannot be known. + """ + fields = schema.fieldDescriptions + fieldID = schema.columnDescriptions[columnID].fFieldID + chain = schema._field_chain(fieldID) + for fid in chain[1:]: + if fields[fid].fStructuralRole in (0x01, 0x03): # collection, variant + msg = ( + f"Column {columnID} is deferred, but its field " + f"{schema.field_path(fieldID)!r} is inside the collection or " + f"variant {schema.field_path(fid)!r}" + ) + raise ValueError(msg) + repetitions = 1 + for fid in chain: + repetitions *= max(fields[fid].fArraySize or 0, 1) + return repetitions + + +class _ClusterReader: + """Builds the clusters of an RNTuple one at a time, from what is computed once""" + + def __init__(self, rntuple: RNTuple): + schema = rntuple.schemaDescription + fieldColumns = schema.field_columns() + self.columns = _columns(schema, fieldColumns) + # The columns of each column's field, by representation and index + self.representations = [ + fieldColumns[column.columnDescription.fFieldID] for column in self.columns + ] + # The elements per entry of each deferred column of the first representation + self.repetitions = [ + _repetitions(schema, column.columnID) + if column.columnDescription.fFirstElementIndex + and column.columnDescription.fRepresentationIndex == 0 + else None + for column in self.columns + ] + self.nHeaderColumns = len(rntuple.headerEnvelope.columnDescriptions.items) + groups = rntuple.footerEnvelope.clusterGroups + if len(rntuple.pagelistEnvelopes) != len(groups): + msg = f"{len(rntuple.pagelistEnvelopes)} page lists for {len(groups)} cluster groups" raise ValueError(msg) - columns = _Columns(schema) - clusters: list[InterpretableCluster] = [] - # The cluster that lists the most columns so far, and how many: only a - # model extension adds columns, so no later cluster lists fewer - widestClusterID, widest = 0, 0 for clusterGroupID, (pagelistEnvelope, group) in enumerate( - zip(self.pagelistEnvelopes, groups, strict=True) + zip(rntuple.pagelistEnvelopes, groups, strict=True) ): - pageLocations = pagelistEnvelope.pageLocations.items - summaries = pagelistEnvelope.clusterSummaries.items - if not len(pageLocations) == len(summaries) == group.fNClusters: + nLocations = len(pagelistEnvelope.pageLocations.items) + nSummaries = len(pagelistEnvelope.clusterSummaries.items) + if not nLocations == nSummaries == group.fNClusters: msg = ( f"Page list of cluster group {clusterGroupID} has " - f"{len(summaries)} cluster summaries and page locations for " - f"{len(pageLocations)} clusters, the group says {group.fNClusters}" + f"{nSummaries} cluster summaries and page locations for " + f"{nLocations} clusters, the group says {group.fNClusters}" ) raise ValueError(msg) - for columnlist, summary in zip(pageLocations, summaries, strict=True): - clusterID = len(clusters) + self.pagelistEnvelopes = rntuple.pagelistEnvelopes + + def clusters(self) -> Iterator[InterpretableCluster]: + nColumns = len(self.columns) + clusterID = 0 + # The cluster that lists the most columns so far, and how many: only a + # model extension adds columns, so no later cluster lists fewer + widestClusterID, widest = 0, 0 + for clusterGroupID, pagelistEnvelope in enumerate(self.pagelistEnvelopes): + for columnlist, summary in zip( + pagelistEnvelope.pageLocations.items, + pagelistEnvelope.clusterSummaries.items, + strict=True, + ): listed = columnlist.items - if not nHeaderColumns <= len(listed) <= nColumns: + if not self.nHeaderColumns <= len(listed) <= nColumns: msg = ( f"Cluster {clusterID} lists {len(listed)} columns; the " - f"schema has {nHeaderColumns} in the header and " + f"schema has {self.nHeaderColumns} in the header and " f"{nColumns} in all" ) raise ValueError(msg) @@ -373,91 +481,21 @@ def clusters(self) -> list[InterpretableCluster]: raise ValueError(msg) if len(listed) > widest: widestClusterID, widest = clusterID, len(listed) - clusters.append( - InterpretableCluster( - clusterID=clusterID, - clusterGroupID=clusterGroupID, - summary=summary, - columns=columns.in_cluster(clusterID, summary, listed), - ) - ) - return clusters - - -@dataclasses.dataclass -class _Column: - """What ``RNTuple.clusters`` needs to know of a column, the same in every cluster""" - - description: ColumnDescription - field: FieldDescription - path: bytes - representations: list[list[int]] - """The columns of the column's field, by representation and index""" - index: int - """The column's index within its representation""" - repetitions: int | None - """For a deferred column of the first representation, its elements per entry""" - - -class _Columns: - """The columns of a schema, and how to find their ranges in a cluster""" - - def __init__(self, schema: SchemaDescription): - fieldColumns = schema.field_columns() - paths: dict[int, bytes] = {} - self.columns: list[_Column] = [] - for columnID, description in enumerate(schema.columnDescriptions): - fieldID = description.fFieldID - if fieldID not in paths: - paths[fieldID] = schema.field_path(fieldID) - representations = fieldColumns[fieldID] - index = representations[description.fRepresentationIndex].index(columnID) - repetitions = None - if description.fFirstElementIndex and description.fRepresentationIndex == 0: - repetitions = self._repetitions(schema, columnID) - self.columns.append( - _Column( - description=description, - field=schema.fieldDescriptions[fieldID], - path=paths[fieldID], - representations=representations, - index=index, - repetitions=repetitions, + yield InterpretableCluster( + clusterID=clusterID, + clusterGroupID=clusterGroupID, + summary=summary, + columnRanges=self._ranges(clusterID, summary, listed), ) - ) - - @staticmethod - def _repetitions(schema: SchemaDescription, columnID: int) -> int: - """The elements per entry of a deferred column of the first representation + clusterID += 1 - The product of the array sizes of its field and its ancestors. The spec - allows an unsuppressed deferred column only where no ancestor is a - collection or a variant (*Column Description*); there, the number of - elements cannot be known. - """ - fields = schema.fieldDescriptions - fieldID = schema.columnDescriptions[columnID].fFieldID - chain = schema._field_chain(fieldID) - for fid in chain[1:]: - if fields[fid].fStructuralRole in (0x01, 0x03): # collection, variant - msg = ( - f"Column {columnID} is deferred, but its field " - f"{schema.field_path(fieldID)!r} is inside the collection or " - f"variant {schema.field_path(fid)!r}" - ) - raise ValueError(msg) - repetitions = 1 - for fid in chain: - repetitions *= max(fields[fid].fArraySize or 0, 1) - return repetitions - - def in_cluster( + def _ranges( self, clusterID: int, summary: ClusterSummary, listed: list[PageLocations[RPageDescription]], - ) -> list[InterpretableColumn]: - """The columns of one cluster, built in the three steps of ROOT's reader""" + ) -> list[InterpretableColumnRange]: + """The column ranges of one cluster, built in the three steps of ROOT's reader""" columns = self.columns locations: list[PageLocations[RPageDescription] | None] = [ *listed, @@ -472,7 +510,9 @@ def in_cluster( zip(locations, columns, strict=True) ): if location is None: - suppressed.append((column.description.fFirstElementIndex or 0) < 0) + suppressed.append( + (column.columnDescription.fFirstElementIndex or 0) < 0 + ) ranges.append((0, 0)) elif location.elementoffset < 0: if location.items: @@ -493,7 +533,7 @@ def in_cluster( for columnID, column in enumerate(columns): if ranges[columnID] is not None: continue - for representation in column.representations: + for representation in self.representations[columnID]: other = representation[column.index] if locations[other] is not None and not suppressed[other]: ranges[columnID] = ranges[other] @@ -501,24 +541,28 @@ def in_cluster( else: msg = ( f"Cluster {clusterID}: column {columnID} is suppressed, and no " - f"other representation of field {column.path!r} is active" + f"other representation of field {column.fieldPath!r} is active" ) raise ValueError(msg) # 3. A deferred column covers the whole cluster, the elements before its # first stored one being zeros; a later representation copies the range # of the first, once that is known - for columnID, column in enumerate(columns): - if column.repetitions is not None: + for columnID, repetitions in enumerate(self.repetitions): + if repetitions is not None: ranges[columnID] = ( - summary.fFirstEntryNumber * column.repetitions, - summary.fNEntries * column.repetitions, + summary.fFirstEntryNumber * repetitions, + summary.fNEntries * repetitions, ) for columnID, column in enumerate(columns): - if column.repetitions is None and column.description.fFirstElementIndex: - ranges[columnID] = ranges[column.representations[0][column.index]] - - out: list[InterpretableColumn] = [] + if ( + self.repetitions[columnID] is None + and column.columnDescription.fFirstElementIndex + ): + first = self.representations[columnID][0][column.index] + ranges[columnID] = ranges[first] + + out: list[InterpretableColumnRange] = [] for columnID, (location, column) in enumerate( zip(locations, columns, strict=True) ): @@ -553,18 +597,15 @@ def in_cluster( firstElementInCluster=nextElement, uncompressedSize=ceil( description.n_elements - * column.description.fBitsOnStorage + * column.columnDescription.fBitsOnStorage / 8 ), # Convert bits to bytes ) ) nextElement += description.n_elements out.append( - InterpretableColumn( - columnID=columnID, - columnDescription=column.description, - fieldDescription=column.field, - fieldPath=column.path, + InterpretableColumnRange( + column=column, pageLocations=location, suppressed=suppressed[columnID], firstElementIndex=firstElementIndex, diff --git a/tests/test_rntuple_clusters.py b/tests/test_rntuple_clusters.py index 8976547..4d196db 100644 --- a/tests/test_rntuple_clusters.py +++ b/tests/test_rntuple_clusters.py @@ -30,13 +30,13 @@ def _ranges(rntuple: RNTuple) -> list[list[tuple[bool, int, int, int, int]]]: return [ [ ( - column.suppressed, - column.firstElementIndex, - column.nElements, - column.nZeroElements, - len(column.pages), + columnRange.suppressed, + columnRange.firstElementIndex, + columnRange.nElements, + columnRange.nZeroElements, + len(columnRange.pages), ) - for column in cluster.columns + for columnRange in cluster.columnRanges ] for cluster in rntuple.clusters() ] @@ -56,18 +56,19 @@ def test_suppressed_columns_keep_their_position(): [(True, 1, 1, 0, 0), (False, 1, 1, 0, 1)], [(False, 2, 1, 0, 1), (True, 2, 1, 0, 0)], ] - cluster = rntuple.clusters()[1] - assert cluster.columns[1].columnDescription.fColumnType == ColumnType.kReal16 - (page,) = cluster.columns[1].pages + cluster = list(rntuple.clusters())[1] + column = cluster.columnRanges[1].column + assert column.columnDescription.fColumnType == ColumnType.kReal16 + (page,) = cluster.columnRanges[1].pages assert page.firstElementInCluster == 0 suppressed = rntuple.pagelistEnvelopes[0].pageLocations[1][0] assert suppressed.elementoffset < 0 - assert cluster.columns[0].pageLocations is suppressed + assert cluster.columnRanges[0].pageLocations is suppressed def test_cluster_ids_continue_across_cluster_groups(): """Issue #116: 3 cluster groups of 5, 4 and 3 clusters are clusters 0 to 11""" - clusters = _load(data_path(MULTIPLE_CLUSTER_GROUPS)).clusters() + clusters = list(_load(data_path(MULTIPLE_CLUSTER_GROUPS)).clusters()) assert [cluster.clusterID for cluster in clusters] == list(range(12)) assert [cluster.clusterGroupID for cluster in clusters] == [0] * 5 + [1] * 4 + [ 2 @@ -78,11 +79,11 @@ def test_cluster_ids_continue_across_cluster_groups(): *(750, 800, 900), ] # Column 2 has two elements per entry; its pages count from its cluster's start - assert [cluster.columns[2].firstElementIndex for cluster in clusters] == [ + assert [cluster.columnRanges[2].firstElementIndex for cluster in clusters] == [ 2 * cluster.summary.fFirstEntryNumber for cluster in clusters ] assert { - cluster.columns[2].pages[0].firstElementInCluster for cluster in clusters + cluster.columnRanges[2].pages[0].firstElementInCluster for cluster in clusters } == {0} @@ -120,11 +121,20 @@ def test_columns_user_class(): ) ) assert len(pinned) == 19 - for cluster in _load(DATA / "user-class.root").clusters(): - assert [ - (column.fieldPath, column.columnDescription.fColumnType) - for column in cluster.columns - ] == [(path, type_) for _, path, type_ in pinned] + rntuple = _load(DATA / "user-class.root") + columns = rntuple.columns() + assert [ + (column.fieldPath, column.columnDescription.fColumnType) for column in columns + ] == [(path, type_) for _, path, type_ in pinned] + # A std::string has an index column and a character column + assert [ + column.index for column in columns if column.fieldPath == b"fHit.fLabel" + ] == [ + 0, + 1, + ] + for cluster in rntuple.clusters(): + assert [columnRange.column for columnRange in cluster.columnRanges] == columns def test_field_paths_schema_extension(): @@ -186,6 +196,37 @@ def test_field_columns_with_a_missing_representation_raise(): rntuple.clusters() +def test_columns_are_described_once(): + """RNTuple.columns() describes each physical column once, with its field and + its index among its representation's columns; every cluster refers to them""" + rntuple = _load(data_path(MULTIPLE_REPRESENTATIONS)) + columns = rntuple.columns() + assert [ + ( + column.columnID, + column.fieldPath, + column.columnDescription.fRepresentationIndex, + column.index, + ) + for column in columns + ] == [(0, b"real", 0, 0), (1, b"real", 1, 0)] + for cluster in rntuple.clusters(): + assert [columnRange.column for columnRange in cluster.columnRanges] == columns + + +def test_clusters_are_built_one_at_a_time(): + """clusters() is an iterator: the schema and the number of page lists and + clusters are checked when it is called, and each cluster's page list when + the iterator reaches it, after the clusters before it were yielded""" + rntuple = _load(data_path(EXTENSION_COLUMNS)) + (pagelist,) = rntuple.pagelistEnvelopes + pagelist.pageLocations.items[2].items.pop() + clusters = rntuple.clusters() + assert [next(clusters).clusterID, next(clusters).clusterID] == [0, 1] + with pytest.raises(ValueError, match="Cluster 2 lists 3 columns"): + next(clusters) + + def test_columns_added_by_model_extension(): """A cluster committed before the model was extended lists only the columns that existed then: here 2 of the 4. The view still has all 4, with the @@ -229,8 +270,11 @@ def test_columns_added_by_model_extension(): (False, 302, 98, 0, 1), ], ] - first = rntuple.clusters()[0] - assert [column.pageLocations is None for column in first.columns] == [ + clusters = list(rntuple.clusters()) + first = clusters[0] + assert [ + columnRange.pageLocations is None for columnRange in first.columnRanges + ] == [ False, False, True, @@ -239,11 +283,11 @@ def test_columns_added_by_model_extension(): # A first stored page starts after the zeros: float_field's at element 200 of # cluster 0, and intvec_field's at element 50 of cluster 1, element 400 of # the column - assert [page.firstElementInCluster for page in first.columns[1].pages] == [200] - column = rntuple.clusters()[1].columns[2] - (page,) = column.pages + assert [page.firstElementInCluster for page in first.columnRanges[1].pages] == [200] + columnRange = clusters[1].columnRanges[2] + (page,) = columnRange.pages assert page.firstElementInCluster == 50 - assert column.firstElementIndex + page.firstElementInCluster == 400 + assert columnRange.firstElementIndex + page.firstElementInCluster == 400 @pytest.mark.parametrize( @@ -275,20 +319,22 @@ def test_every_column_is_complete(path: Path | str): rntuple = _load(data_path(path)) nColumns = len(rntuple.schemaDescription.columnDescriptions) nextElement = [0] * nColumns + columns = rntuple.columns() + assert [column.columnID for column in columns] == list(range(nColumns)) for cluster in rntuple.clusters(): - assert [column.columnID for column in cluster.columns] == list(range(nColumns)) - for column in cluster.columns: - if column.suppressed: - assert column.pages == [] - assert column.nZeroElements == 0 + assert [columnRange.column for columnRange in cluster.columnRanges] == columns + for columnID, columnRange in enumerate(cluster.columnRanges): + if columnRange.suppressed: + assert columnRange.pages == [] + assert columnRange.nZeroElements == 0 else: - start = column.nZeroElements - for page in column.pages: + start = columnRange.nZeroElements + for page in columnRange.pages: assert page.firstElementInCluster == start start += page.pageDescription.n_elements - assert start == column.nElements - assert column.firstElementIndex == nextElement[column.columnID] - nextElement[column.columnID] += column.nElements + assert start == columnRange.nElements + assert columnRange.firstElementIndex == nextElement[columnID] + nextElement[columnID] += columnRange.nElements def test_page_list_with_a_missing_header_column_raises(): @@ -297,7 +343,7 @@ def test_page_list_with_a_missing_header_column_raises(): (pagelist,) = rntuple.pagelistEnvelopes pagelist.pageLocations.items[0].items = [] with pytest.raises(ValueError, match="Cluster 0 lists 0 columns"): - rntuple.clusters() + list(rntuple.clusters()) def test_page_list_with_an_extra_column_raises(): @@ -305,7 +351,7 @@ def test_page_list_with_an_extra_column_raises(): columnlist = rntuple.pagelistEnvelopes[1].pageLocations.items[2] columnlist.items.append(columnlist.items[0]) with pytest.raises(ValueError, match="Cluster 7 lists 4 columns"): - rntuple.clusters() + list(rntuple.clusters()) def test_page_list_with_fewer_columns_than_an_earlier_one_raises(): @@ -316,7 +362,7 @@ def test_page_list_with_fewer_columns_than_an_earlier_one_raises(): with pytest.raises( ValueError, match="Cluster 2 lists 3 columns, fewer than the 4 of cluster 1" ): - rntuple.clusters() + list(rntuple.clusters()) def test_page_lists_must_match_the_cluster_groups(): @@ -350,7 +396,7 @@ def test_suppressed_column_with_pages_raises(): clusters = rntuple.pagelistEnvelopes[0].pageLocations.items clusters[1].items[0].items = list(clusters[0].items[0].items) with pytest.raises(ValueError, match="column 0 is suppressed but has 1 pages"): - rntuple.clusters() + list(rntuple.clusters()) def test_field_with_every_representation_suppressed_raises(): @@ -359,7 +405,7 @@ def test_field_with_every_representation_suppressed_raises(): clusters = rntuple.pagelistEnvelopes[0].pageLocations.items clusters[1].items[1] = clusters[1].items[0] with pytest.raises(ValueError, match="no other representation of field b'real'"): - rntuple.clusters() + list(rntuple.clusters()) def test_deferred_column_whose_pages_disagree_raises(): @@ -370,7 +416,7 @@ def test_deferred_column_whose_pages_disagree_raises(): with pytest.raises( ValueError, match="column 2 has elements 350 to 467, but its pages hold 67" ): - rntuple.clusters() + list(rntuple.clusters()) def test_deferred_column_inside_a_collection_raises(): diff --git a/tests/test_rntuple_hardcoded.py b/tests/test_rntuple_hardcoded.py index b02d4ab..31141cc 100644 --- a/tests/test_rntuple_hardcoded.py +++ b/tests/test_rntuple_hardcoded.py @@ -338,9 +338,10 @@ def test_read_contributors(): extraTypeInformations=[], ) - clusters = rntuple.clusters() + clusters = list(rntuple.clusters()) assert [ - [column.pages for column in cluster.columns] for cluster in clusters + [columnRange.pages for columnRange in cluster.columnRanges] + for cluster in clusters ] == [ # Clusters (columnlists) [ # Columns (pagelists) [ # Pages (page descriptions) @@ -387,7 +388,7 @@ def test_read_contributors(): ] assert [ (column.columnDescription.fColumnType, column.fieldPath) - for column in clusters[0].columns + for column in rntuple.columns() ] == [ (ColumnType.kIndex64, b"firstName"), (ColumnType.kChar, b"firstName"), @@ -595,9 +596,10 @@ def fetch_from_locator(loc): extraTypeInformations=[], ) - clusters_a = rntuple_a.clusters() + clusters_a = list(rntuple_a.clusters()) assert [ - [column.pages for column in cluster.columns] for cluster in clusters_a + [columnRange.pages for columnRange in cluster.columnRanges] + for cluster in clusters_a ] == [ # Clusters (columnlists) [ # Columns (pagelists) [ # Pages (page descriptions) @@ -614,7 +616,7 @@ def fetch_from_locator(loc): ] assert [ (column.columnDescription.fColumnType, column.fieldPath) - for column in clusters_a[0].columns + for column in rntuple_a.columns() ] == [(ColumnType.kSplitReal32, b"f")] anchor_b = keylist[b"B"].read_object(fetch_data, ROOT3a3aRNTuple) @@ -792,9 +794,10 @@ def fetch_from_locator(loc): extraTypeInformations=[], ) - clusters_b = rntuple_b.clusters() + clusters_b = list(rntuple_b.clusters()) assert [ - [column.pages for column in cluster.columns] for cluster in clusters_b + [columnRange.pages for columnRange in cluster.columnRanges] + for cluster in clusters_b ] == [ # Clusters (columnlists) [ # Columns (pagelists) [ # Pages (page descriptions) @@ -811,5 +814,5 @@ def fetch_from_locator(loc): ] assert [ (column.columnDescription.fColumnType, column.fieldPath) - for column in clusters_b[0].columns + for column in rntuple_b.columns() ] == [(ColumnType.kSplitInt32, b"g")]