diff --git a/src/rootfilespec/rntuple/RNTuple.py b/src/rootfilespec/rntuple/RNTuple.py index a5c9aa6..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 @@ -10,12 +10,11 @@ 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, - ColumnType, ExtraTypeInformation, FieldDescription, ) @@ -79,19 +78,187 @@ def from_envelopes( extraTypeInformations, ) + 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. Raises on an out-of-range ID or a cycle. + """ + fields = self.fieldDescriptions + 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 chain: + msg = ( + f"Field {field_id}: the parent chain has a cycle at field ID {fid}" + ) + raise ValueError(msg) + 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]) + + 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 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: - """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 start in the cluster. """ 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.""" + 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 ``InterpretableColumnRange.firstElementIndex`` plus this.""" 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.""" + """The size of the page's elements packed for storage, before compression, in bytes.""" + + +@dataclasses.dataclass +class InterpretableColumnRange: + """One physical column in one cluster: its element range and its pages + + 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). + """ + + 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. + + ``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 InterpretableCluster: + """One cluster: its entries, and every column's elements and pages in it + + Like ROOT's ``RClusterDescriptor``. ``clusterID``, + ``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*). + """ + + 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.""" + columnRanges: list[InterpretableColumnRange] + """One per physical column, in column-ID order, as the spec guarantees + (*Page Locations*): a position is the column ID.""" @dataclasses.dataclass @@ -173,40 +340,278 @@ 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. - 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. + def columns(self) -> list[InterpretableColumn]: + """Every physical column, in column-ID order, with its field and its place among the field's columns - Args: - includeSuppressed (bool): If False, skip suppressed 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``). """ - 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 + 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 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. 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. + """ + 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) + for clusterGroupID, (pagelistEnvelope, group) in enumerate( + zip(rntuple.pagelistEnvelopes, groups, strict=True) + ): + 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"{nSummaries} cluster summaries and page locations for " + f"{nLocations} clusters, the group says {group.fNClusters}" + ) + raise ValueError(msg) + 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 self.nHeaderColumns <= len(listed) <= nColumns: + msg = ( + f"Cluster {clusterID} lists {len(listed)} columns; the " + f"schema has {self.nHeaderColumns} in the header and " + f"{nColumns} in all" ) - if includeSuppressed or pagelist.elementoffset >= 0 - ] - for columnlist in pagelistEnvelope.pageLocations - ] - for pagelistEnvelope in self.pagelistEnvelopes + 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) + yield InterpretableCluster( + clusterID=clusterID, + clusterGroupID=clusterGroupID, + summary=summary, + columnRanges=self._ranges(clusterID, summary, listed), + ) + clusterID += 1 + + def _ranges( + self, + clusterID: int, + summary: ClusterSummary, + listed: list[PageLocations[RPageDescription]], + ) -> 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, + *[None] * (len(columns) - len(listed)), ] - return envelopePages + + # 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.columnDescription.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) + 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 representation in self.representations[columnID]: + 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 {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, repetitions in enumerate(self.repetitions): + if repetitions is not None: + ranges[columnID] = ( + summary.fFirstEntryNumber * repetitions, + summary.fNEntries * repetitions, + ) + for columnID, column in enumerate(columns): + 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) + ): + 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}" + ) + raise ValueError(msg) + elif not suppressed[columnID]: + nZeroElements = nElements + pages: list[InterpretablePage] = [] + nextElement = nZeroElements + for description in descriptions: + pages.append( + InterpretablePage( + pageDescription=description, + firstElementInCluster=nextElement, + uncompressedSize=ceil( + description.n_elements + * column.columnDescription.fBitsOnStorage + / 8 + ), # Convert bits to bytes + ) + ) + nextElement += description.n_elements + out.append( + InterpretableColumnRange( + column=column, + pageLocations=location, + suppressed=suppressed[columnID], + firstElementIndex=firstElementIndex, + nElements=nElements, + nZeroElements=nZeroElements, + pages=pages, + ) + ) + return out diff --git a/tests/test_rntuple_clusters.py b/tests/test_rntuple_clusters.py new file mode 100644 index 0000000..4d196db --- /dev/null +++ b/tests/test_rntuple_clusters.py @@ -0,0 +1,442 @@ +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 + +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: + with open_path(path) as reader: + keylist = reader.keylist() + (name,) = [n for n in keylist if keylist[n].fClassName == b"ROOT::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 [ + [ + ( + columnRange.suppressed, + columnRange.firstElementIndex, + columnRange.nElements, + columnRange.nZeroElements, + len(columnRange.pages), + ) + for columnRange in cluster.columnRanges + ] + for cluster in rntuple.clusters() + ] + + +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. A suppressed column + takes the element range of the active one (spec, *Suppressed Columns*). + """ + 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)], + ] + 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.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 = 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 + ] * 3 + assert [cluster.summary.fFirstEntryNumber for cluster in clusters] == [ + *(0, 100, 200, 300, 400), + *(450, 500, 600, 700), + *(750, 800, 900), + ] + # Column 2 has two elements per entry; its pages count from its cluster's start + assert [cluster.columnRanges[2].firstElementIndex for cluster in clusters] == [ + 2 * cluster.summary.fFirstEntryNumber for cluster in clusters + ] + assert { + cluster.columnRanges[2].pages[0].firstElementInCluster for cluster in clusters + } == {0} + + +@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] + + +@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") + 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(): + """Field IDs continue from the header into the footer's schema extension""" + 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)] == [ + 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") +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_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 + extension[0] = dataclasses.replace(extension[0], fRepresentationIndex=1) + with pytest.raises(ValueError, match=r"Field 1 has representations \[1\]"): + 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 + 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), + ], + ] + clusters = list(rntuple.clusters()) + first = clusters[0] + assert [ + columnRange.pageLocations is None for columnRange in first.columnRanges + ] == [ + False, + False, + True, + True, + ] + # 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.columnRanges[1].pages] == [200] + columnRange = clusters[1].columnRanges[2] + (page,) = columnRange.pages + assert page.firstElementInCluster == 50 + assert columnRange.firstElementIndex + page.firstElementInCluster == 400 + + +@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 [ + MULTIPLE_REPRESENTATIONS, + MULTIPLE_CLUSTER_GROUPS, + EXTENSION_COLUMNS, + ] + ), + ], +) +def test_every_column_is_complete(path: Path | str): + """In every cluster: one entry per column; a suppressed column has no pages; + 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)) + 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 [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 = columnRange.nZeroElements + for page in columnRange.pages: + assert page.firstElementInCluster == start + start += page.pageDescription.n_elements + assert start == columnRange.nElements + assert columnRange.firstElementIndex == nextElement[columnID] + nextElement[columnID] += columnRange.nElements + + +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"): + list(rntuple.clusters()) + + +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"): + list(rntuple.clusters()) + + +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" + ): + list(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.clusters() + + +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.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.clusters() + + +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"): + list(rntuple.clusters()) + + +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'"): + list(rntuple.clusters()) + + +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" + ): + list(rntuple.clusters()) + + +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.clusters() + + +def test_field_path_errors(): + 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) + # _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) diff --git a/tests/test_rntuple_hardcoded.py b/tests/test_rntuple_hardcoded.py index 1b743af..31141cc 100644 --- a/tests/test_rntuple_hardcoded.py +++ b/tests/test_rntuple_hardcoded.py @@ -338,53 +338,63 @@ def test_read_contributors(): extraTypeInformations=[], ) - extended_page_descriptions = rntuple.get_extended_page_descriptions() - assert 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, - ) - ], - [ - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-178, - locator=StandardLocator(size=178, offset=804), - ), - uncompressedSize=178, - columnType=ColumnType.kChar, - ) - ], - [ - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-22, - locator=StandardLocator(size=176, offset=990), - ), - uncompressedSize=176, - columnType=ColumnType.kIndex64, - ) - ], - [ - InterpretablePage( - pageDescription=RPageDescription( - fNElements=-193, - locator=StandardLocator(size=193, offset=1174), - ), - uncompressedSize=193, - columnType=ColumnType.kChar, - ) - ], - ] + clusters = list(rntuple.clusters()) + assert [ + [columnRange.pages for columnRange in cluster.columnRanges] + for cluster in clusters + ] == [ # Clusters (columnlists) + [ # Columns (pagelists) + [ # Pages (page descriptions) + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-22, + locator=StandardLocator(size=176, offset=620), + ), + firstElementInCluster=0, + uncompressedSize=176, + ) + ], + [ + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-178, + locator=StandardLocator(size=178, offset=804), + ), + firstElementInCluster=0, + uncompressedSize=178, + ) + ], + [ + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-22, + locator=StandardLocator(size=176, offset=990), + ), + firstElementInCluster=0, + uncompressedSize=176, + ) + ], + [ + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-193, + locator=StandardLocator(size=193, offset=1174), + ), + firstElementInCluster=0, + uncompressedSize=193, + ) + ], ] ] + assert [ + (column.columnDescription.fColumnType, column.fieldPath) + for column in rntuple.columns() + ] == [ + (ColumnType.kIndex64, b"firstName"), + (ColumnType.kChar, b"firstName"), + (ColumnType.kIndex64, b"lastName"), + (ColumnType.kChar, b"lastName"), + ] def test_read_multiple_rntuples(): @@ -586,23 +596,28 @@ def fetch_from_locator(loc): extraTypeInformations=[], ) - extended_page_descriptions_a = rntuple_a.get_extended_page_descriptions() - assert 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, - ) - ] + clusters_a = list(rntuple_a.clusters()) + assert [ + [columnRange.pages for columnRange in cluster.columnRanges] + for cluster in clusters_a + ] == [ # Clusters (columnlists) + [ # Columns (pagelists) + [ # Pages (page descriptions) + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-100, + locator=StandardLocator(size=138, offset=409), + ), + firstElementInCluster=0, + uncompressedSize=400, + ) ] ] ] + assert [ + (column.columnDescription.fColumnType, column.fieldPath) + for column in rntuple_a.columns() + ] == [(ColumnType.kSplitReal32, b"f")] anchor_b = keylist[b"B"].read_object(fetch_data, ROOT3a3aRNTuple) assert anchor_b == ROOT3a3aRNTuple( @@ -779,20 +794,25 @@ def fetch_from_locator(loc): extraTypeInformations=[], ) - extended_page_descriptions_b = rntuple_b.get_extended_page_descriptions() - assert 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, - ) - ] + clusters_b = list(rntuple_b.clusters()) + assert [ + [columnRange.pages for columnRange in cluster.columnRanges] + for cluster in clusters_b + ] == [ # Clusters (columnlists) + [ # Columns (pagelists) + [ # Pages (page descriptions) + InterpretablePage( + pageDescription=RPageDescription( + fNElements=-100, + locator=StandardLocator(size=164, offset=1695), + ), + firstElementInCluster=0, + uncompressedSize=400, + ) ] ] ] + assert [ + (column.columnDescription.fColumnType, column.fieldPath) + for column in rntuple_b.columns() + ] == [(ColumnType.kSplitInt32, b"g")]