diff --git a/src/rootfilespec/rntuple/RPage.py b/src/rootfilespec/rntuple/RPage.py index 6e45576..3f6c6db 100644 --- a/src/rootfilespec/rntuple/RPage.py +++ b/src/rootfilespec/rntuple/RPage.py @@ -8,7 +8,9 @@ class RPage(ROOTSerializable): """A class to represent an RNTuple page.""" page: bytes - """The RNTuple page raw data.""" + """The RNTuple page raw data, as stored (still compressed if the page is).""" + checksum: int | None = None + """The page's XXH3-64 checksum, verified against ``page``, or None if it has none.""" # TODO: Flush out RPage class @classmethod diff --git a/src/rootfilespec/rntuple/pagelocations.py b/src/rootfilespec/rntuple/pagelocations.py index 677b8f1..dc1e7dd 100644 --- a/src/rootfilespec/rntuple/pagelocations.py +++ b/src/rootfilespec/rntuple/pagelocations.py @@ -1,6 +1,8 @@ from collections.abc import Callable from typing import Annotated, cast +import xxhash # type: ignore[import-not-found] + from rootfilespec.bootstrap.compression import RCompressionSettings from rootfilespec.rntuple.RFrame import Item, ListFrame from rootfilespec.rntuple.RLocator import RLocator @@ -13,6 +15,9 @@ ) from rootfilespec.structutil import Fmt, OptionalField +PAGE_CHECKSUM_SIZE = 8 +"""A page checksum is an XXH3-64, stored little-endian right after the page""" + @serializable class RPageDescription(ROOTSerializable): @@ -30,7 +35,11 @@ class RPageDescription(ROOTSerializable): """ fNElements: Annotated[int, Fmt(" int: @property def size(self) -> int: - """The (compressed) size of the page data.""" - return self.locator.size + """The number of bytes to fetch: the page as stored, then its checksum if it has one. + + The spec's page size, which excludes the checksum, is ``locator.size``.""" + return self.locator.size + (PAGE_CHECKSUM_SIZE if self.has_checksum else 0) + + @property + def n_elements(self) -> int: + """The number of elements in the page.""" + return abs(self.fNElements) + + @property + def has_checksum(self) -> bool: + """Whether an XXH3-64 checksum is stored right after the page.""" + return self.fNElements < 0 def read_from(self, buffer: ReadBuffer) -> RPage: - """Read the page from the given buffer. + """Read the page from the given buffer, and verify its checksum if it has one. - Pages are wrapped in compression blocks (like envelopes). + The buffer holds ``size`` bytes. The checksum covers the page as stored + (sealed, possibly compressed), wherever those bytes were fetched from. + Pages are wrapped in compression blocks (like envelopes); nothing is + decompressed here. """ - #### Read the page from the buffer - page, buffer = RPage.read(buffer) - - if buffer: - msg = "RPageDescription.read_from: buffer not empty after reading page." + if len(buffer) != self.size: + msg = f"RPageDescription.read_from: expected {self.size} bytes, got {len(buffer)}" raise ValueError(msg) - return page + #### Read the page from the buffer + if not self.has_checksum: + page, _ = RPage.read(buffer) + return page + + data, buffer = buffer.consume(self.locator.size) + (checksum,), buffer = buffer.unpack(" list[RPageDescription]: + out: list[RPageDescription] = [] + with open_path(path) as reader: + keylist = reader.keylist() + for name in keylist: + key = keylist[name] + if key.fClassName.fString != b"ROOT::RNTuple": + continue + rntuple = RNTuple.from_anchor(reader.fetch(key), reader.fetch.buffer) + for pagelist in rntuple.pagelistEnvelopes: + for cluster in pagelist.pageLocations: + for column in cluster: + out.extend(column) + return out + + +def _buffer(raw: bytes, offset: int, size: int) -> ReadBuffer: + return ReadBuffer( + memoryview(raw[offset : offset + size]), + 0, + BOOTSTRAP_CONTEXT, + BufferContext(abspos=offset), + ) + + +def test_page_with_checksum(): + """Issue #55: the page at 550 in rntuple/anchor.root, checked against its bytes + + Its fNElements is -3 (the sign flags a checksum), its locator says 12 bytes, + and the XXH3-64 of those 12 bytes is stored little-endian at 562..570. The + page description's size covers both, so one fetch gets what read_from needs. + """ + path = DATA / "anchor.root" + raw = path.read_bytes() + (page,) = [p for p in _page_descriptions(path) if p.offset == 550] + assert page.fNElements == -3 + assert page.n_elements == 3 + assert page.has_checksum + assert page.locator.size == 12 + assert page.size == 20 + + assert raw[562:570].hex() == "de3ce2c4a5a407be" + read = page.read_from(_buffer(raw, page.offset, page.size)) + assert read.page == raw[550:562] + assert read.checksum == int.from_bytes(raw[562:570], "little") + assert read.checksum == xxhash.xxh3_64_intdigest(raw[550:562]) + + +def test_corrupted_page_raises(): + path = DATA / "anchor.root" + raw = bytearray(path.read_bytes()) + (page,) = [p for p in _page_descriptions(path) if p.offset == 550] + raw[555] ^= 0x01 + with pytest.raises(ValueError, match=r"Page checksum mismatch at .*offset=550"): + page.read_from(_buffer(bytes(raw), page.offset, page.size)) + + +def test_wrong_length_raises(): + page = RPageDescription(-3, StandardLocator(12, 0)) + with pytest.raises(ValueError, match="expected 20 bytes"): + page.read_from(_buffer(bytes(12), 0, 12)) + + +def test_page_without_checksum(): + raw = b"0123456789" + page = RPageDescription(3, StandardLocator(10, 0)) + assert not page.has_checksum + assert page.size == 10 + read = page.read_from(_buffer(raw, 0, 10)) + assert read.page == raw + assert read.checksum is None + + +@pytest.mark.parametrize("name", sorted(p.name for p in DATA.glob("*.root"))) +def test_every_page_verifies(name: str): + """Every page of every root-io-spec RNTuple fixture passes its checksum + + Pages are read by their stored bytes, never decompressed, and several page + descriptions may name the same bytes (same-page merging): each still reads. + """ + path = DATA / name + raw = path.read_bytes() + pages = _page_descriptions(path) + assert pages + for page in pages: + assert page.offset + page.size <= len(raw) + read = page.read_from(_buffer(raw, page.offset, page.size)) + assert read.checksum is not None + assert len(read.page) == page.locator.size + + +def test_shared_page_ranges(): + """rntuple/map.root has page descriptions that name the same bytes""" + ranges = [(p.offset, p.size) for p in _page_descriptions(DATA / "map.root")] + assert len(set(ranges)) < len(ranges)