Parallel b-tree reading

[bnlawrence]

Parallel B-Tree Reading via cat_ranges

Much usage of pyfive will be for remote access to data where the data is presented to pyfive
vai fsspec, which itself supports a significant set of benefits, the most of important of
which is that if the requests are provided together via cat_ranges, rather than a series of requests,
each with latency, the requests are generated concurrently. Tests with chunk reading show this
can improve performance by a factor of 10 on remote http and s3 resources.
Some benefit ought to be possible for b-trees as well.

Background and Motivation

HDF5 B-tree traversal is currently sequential: each node is read via fh.seek() +
fh.read(), and child addresses are only known after the parent node has been parsed.
For remote files this means O(N nodes) serial round-trips.

The existing MetadataBufferingWrapper papers over this by eagerly reading the first
1 MiB on first access, which covers B-tree nodes for well-packed files. For
poorly-packed files (nodes scattered beyond 1 MiB) the fallback is individual remote
reads — exactly the pattern that cat_ranges is designed to replace.

The chunk data path already uses cat_ranges via _read_bulk_fsspec. This change
applies the same technique to B-tree node reads.

---

Key Insight: What Can Be Parallelised

B-tree traversal has a sequential dependency between levels: child addresses are
only discovered by parsing the parent. This cannot be avoided. However, once all
addresses at a given level are known, every read at that level is independent.

Level N (root):   sequential read          ← 1 read, unavoidable
Level N-1:        cat_ranges               ← all addresses known after root
...
Level 0 (leaves): cat_ranges               ← all addresses known after level 1

In practice HDF5 v1 B-trees are shallow (2–3 levels), so the internal nodes are few
and cheap. The leaves dominate, and that is where the win is largest.

The target optimisation: traverse internal nodes sequentially (cheap, few reads),
collect all leaf addresses, then fetch all leaf nodes in a single cat_ranges call.

---

Scope

This change applies to BTreeV1RawDataChunks only — this is the B-tree used for
chunked raw data and is the hot path for array reads. BTreeV1Groups and the v2
B-trees are lower priority and should not be changed in this pass.

---

Implementation

1. Add a _supports_cat_ranges helper on AbstractBTree

@staticmethod
def _get_cat_ranges_fs(fh):
    """
    Return (fs, path) if fh supports cat_ranges, else (None, None).
    Reaches through MetadataBufferingWrapper the same way _read_bulk_fsspec does.
    """
    actual_fh = getattr(fh, "fh", fh)
    fs = getattr(actual_fh, "fs", None)
    path = getattr(actual_fh, "path", None)
    if fs is not None and path is not None and hasattr(fs, "cat_ranges"):
        return fs, path
    return None, None

2. Split BTreeV1RawDataChunks.__init__ traversal into two phases

The existing _read_children walks the tree level-by-level sequentially. Replace it
in BTreeV1RawDataChunks with a version that:

• Phase 1: reads internal nodes (level > 0) sequentially as today
• Phase 2: collects all leaf addresses from the parsed internal nodes, then fetches
  all leaf nodes in one cat_ranges call

def _read_children(self):
    """Override to parallelise leaf reads via cat_ranges when available."""
    fs, path = self._get_cat_ranges_fs(self.fh)

    if fs is None:
        # No cat_ranges support — fall back to original sequential traversal
        super()._read_children()
        return

    # Phase 1: traverse internal nodes sequentially (level > 0)
    # These are few in number; sequential reads are fine here, and we
    # cannot avoid the dependency chain between levels.
    for node_level in range(self.depth, 0, -1):
        for parent_node in self.all_nodes[node_level]:
            for child_addr in parent_node["addresses"]:
                if node_level - 1 > 0:
                    # Still an internal node — read sequentially as before
                    child_node = self._read_node(child_addr, node_level - 1)
                    self._add_node(child_node)
                # If node_level - 1 == 0 it's a leaf; collect below

    # Phase 2: collect all leaf addresses from the lowest internal level
    leaf_addresses = []
    lowest_internal = self.all_nodes.get(1, [])
    for node in lowest_internal:
        leaf_addresses.extend(node["addresses"])

    if not leaf_addresses:
        # Depth-0 tree (root is also leaf) — already read, nothing to do
        return

    # Phase 3: bulk-fetch all leaf nodes in one cat_ranges call
    leaf_size = self._estimate_leaf_node_size()
    starts = [addr for addr in leaf_addresses]
    stops  = [addr + leaf_size for addr in leaf_addresses]

    raw_leaves = fs.cat_ranges([path] * len(leaf_addresses), starts, stops)

    for addr, raw in zip(leaf_addresses, raw_leaves):
        node = self._parse_leaf_from_buffer(raw, addr)
        self._add_node(node)

3. Leaf node size estimation

HDF5 v1 B-tree node size is not stored in the node itself; it must be inferred.
The safest approach is to read the first leaf node sequentially to determine its
actual size, then use that for all subsequent reads.

def _estimate_leaf_node_size(self):
    """
    Return the byte size of a leaf node by reading the first one sequentially.

    HDF5 v1 B-tree nodes have no stored size field; the size is a function of
    the number of entries (entries_used in the header) and the per-entry record
    size (which depends on self.dims).  We read one node to get entries_used,
    compute the size, and assume all leaf nodes are the same size.

    Per-entry layout for NODE_TYPE=1 (raw data chunks):
        8 bytes  : chunk_size + filter_mask (two uint32)
        8 * dims : chunk_offset (dims uint64 values)
        8 bytes  : chunk_address (uint64)
    """
    # B_LINK_NODE header size: 4 + 1 + 1 + 2 + 8 + 8 = 24 bytes
    HEADER_SIZE = 24
    # Per-entry record size
    entry_size = 8 + 8 * self.dims + 8
    # N+1 keys are *not* present in NODE_TYPE=1 (unlike type 0 groups)

    # Peek at the first leaf node to get entries_used
    first_leaf_addr = self.all_nodes[1][0]["addresses"][0]
    self.fh.seek(first_leaf_addr)
    header_bytes = self.fh.read(HEADER_SIZE)
    entries_used = struct.unpack_from("<H", header_bytes, 6)[0]  # offset 6 in header

    return HEADER_SIZE + entries_used * entry_size

Note: If leaf nodes can have differing entries_used values (e.g. the last leaf
in a non-full tree), then using a fixed size will over-read the last few nodes.
Over-reading is harmless — _parse_leaf_from_buffer reads only entries_used
entries from the buffer — but wastes a small amount of bandwidth. An alternative
is to compute the size per-address by peeking all headers in a first cat_ranges
call (24 bytes each), then issuing a second cat_ranges for full node bodies.
For most files the single-size approach is sufficient.

4. Parse a leaf node from a raw buffer

Extract the existing sequential parse logic from _read_node into a buffer-based
variant so it can be called on the raw bytes returned by cat_ranges:

def _parse_leaf_from_buffer(self, raw, addr):
    """
    Parse a leaf node (node_level=0) from a raw bytes buffer.
    Mirrors the logic in _read_node / _read_node_header but reads from
    a BytesIO rather than seeking self.fh.
    """
    import io
    buf = io.BytesIO(raw)

    # --- header (mirrors _read_node_header) ---
    node = _unpack_struct_from(self.B_LINK_NODE, buf.read(struct.calcsize(
        "<" + "".join(self.B_LINK_NODE.values())
    )))
    assert node["signature"] == b"TREE"
    assert node["node_type"] == self.NODE_TYPE
    assert node["node_level"] == 0

    # --- entries (mirrors _read_node body for NODE_TYPE=1) ---
    keys = []
    addresses = []
    fmt = "<" + "Q" * self.dims
    fmt_size = struct.calcsize(fmt)

    for _ in range(node["entries_used"]):
        chunk_size, filter_mask = struct.unpack("<II", buf.read(8))
        chunk_offset = struct.unpack(fmt, buf.read(fmt_size))
        chunk_address = struct.unpack("<Q", buf.read(8))[0]
        keys.append(OrderedDict((
            ("chunk_size",   chunk_size),
            ("filter_mask",  filter_mask),
            ("chunk_offset", chunk_offset),
        )))
        addresses.append(chunk_address)

    node["keys"] = keys
    node["addresses"] = addresses
    self.last_offset = max(addr, self.last_offset)
    return node

Implementation note: _unpack_struct_from already exists in core.py and
accepts a buffer argument — use it directly rather than duplicating the unpacking
logic. Check its exact signature before wiring up.

---

Interaction with MetadataBufferingWrapper

cat_ranges is called on actual_fh.fs (the underlying fsspec filesystem), bypassing
MetadataBufferingWrapper entirely — exactly as _read_bulk_fsspec does for chunks.
This is correct behaviour:

• Well-packed files: the 1 MiB eager buffer already holds the leaf nodes.
  cat_ranges will hit the fsspec cache layer and cost nothing extra.
• Poorly-packed files: cat_ranges issues parallel remote fetches for all
  scattered leaf nodes in one round-trip, which is far better than the current
  sequential fallback reads.

No changes to MetadataBufferingWrapper are required.

---

Fallback Behaviour

_read_children checks for cat_ranges support before doing anything new. If the
filesystem does not support cat_ranges (local files, in-memory buffers, custom
wrappers) the call falls through to super()._read_children() — the existing
sequential traversal — unchanged.

---

Testing

1. Unit test — well-packed file on S3 (mocked): mock fs.cat_ranges to return
   pre-built leaf node bytes; assert it is called once with all leaf addresses; assert
   the resulting self.all_nodes[0] matches the sequential parse.

2. Unit test — no cat_ranges: pass a plain BytesIO fh; assert the sequential
   path is taken and results are identical.

3. Regression: run the existing B-tree test suite unchanged — all results must
   match before and after.

4. Integration: open a real chunked HDF5 file via s3fs; assert chunk index
   contents are identical between serial and parallel paths.

---

Files to Change

File	Change
pyfive/btree.py	Add _get_cat_ranges_fs, override _read_children and add _estimate_leaf_node_size / _parse_leaf_from_buffer to BTreeV1RawDataChunks
tests/test_btree_parallel.py	New test file covering the cases above

No changes required to metadata_buffering_wrapper.py, chunk_read.py, or any caller
of BTreeV1RawDataChunks — the parallel path is entirely internal to the class.

[davidhassell]

Hi Bryan, a lot to absorb here ...

Phase 1: reads internal nodes (level > 0) sequentially as today

so what's the benefit? I'm sure I'm missing something!

[bnlawrence]

I think that's the language of a b-tree, the internal nodes are the non-leaf nodes — they don't contain any chunk data themselves, they just contain the addresses of their children. We have to read them to get the leaves which is where we get the parallelism.

[bnlawrence]

First tests with a non-trivial use case suggest a factor of two in performance, but the amount will depend on the complexity of the b-tree.

[bnlawrence]

After implementing that, and investigating it, and thinking about the application to POSIX as well (where some b-trees are pathological there too), a slightly different design emerges:

Background

HDF5 v1 B-tree traversal is currently sequential: every node is fetched via
fh.seek() + fh.read(). For files with many chunks (e.g. 292,192 chunks in
uas_3hr_IPSL-CM6A-LR_...), the B-tree leaf nodes are scattered across the entire
file and the sequential traversal dominates read time even on POSIX.

This spec describes adding parallel leaf-node fetching via a fetch_fn callable
injected into BTreeV1RawDataChunks at construction time. The BTree class itself
remains unaware of any parallelism strategy; the caller (DatasetID) supplies a
fetch_fn closure built from its existing parallelism configuration.

This design was chosen over alternatives because:

• Parallelism configuration is already owned by DatasetID via set_parallelism.
• The same unified strategy (cat_ranges, os.pread threads, or serial) can be
  applied consistently to both chunk reads and B-tree leaf reads.
• As parallelism moves higher up the stack and becomes opt-out, fetch_fn is simply
  constructed at that higher level and passed down — nothing in BTreeV1RawDataChunks
  needs to change.

---

Scope

Changes are confined to BTreeV1RawDataChunks (leaf fetch parallelism) and
DatasetID (constructing and passing fetch_fn). AbstractBTree, BTreeV1,
BTreeV1Groups, BTreeV2, and MetadataBufferingWrapper are untouched.

---

Why Internal Nodes Remain Sequential

Internal nodes (level > 0) contain only child addresses — no chunk data. Their
addresses are discovered by reading their parent, creating an unavoidable sequential
dependency chain from the root downwards. For real files internal nodes are few
(single digits even for 292,192-chunk datasets) so sequential reads are negligible.
Only the leaf nodes (level 0) are numerous enough to warrant parallel fetching.

---

The fetch_fn Interface

fetch_fn(addresses: list[int], size: int) -> list[bytes]

• addresses: list of file offsets, one per leaf node
• size: byte size to read at each address (all leaf nodes are the same size —
  see Leaf Node Size below)
• returns: list of raw byte buffers in the same order as addresses

fetch_fn is optional. If None (the default), BTreeV1RawDataChunks falls back
to the existing sequential fh.seek() / fh.read() traversal unchanged.

---

Changes to BTreeV1RawDataChunks

1. Accept fetch_fn in __init__

def __init__(self, fh, offset, dims, fetch_fn=None):
    self.dims = dims
    self._fetch_fn = fetch_fn
    super().__init__(fh, offset)

2. Override _read_children

Replace the inherited sequential traversal with a two-phase version when fetch_fn
is provided:

def _read_children(self):
    """
    Phase 1: traverse internal nodes sequentially (unavoidable dependency chain).
    Phase 2: fetch all leaf nodes in parallel via fetch_fn (if provided).
    """
    if self._fetch_fn is None:
        super()._read_children()
        return

    # Phase 1 — internal nodes (level > 0), sequential as before.
    # For typical files this is O(1)–O(tens) reads.
    for node_level in range(self.depth, 0, -1):
        for parent_node in self.all_nodes[node_level]:
            for child_addr in parent_node["addresses"]:
                if node_level - 1 > 0:
                    child_node = self._read_node(child_addr, node_level - 1)
                    self._add_node(child_node)
                # level-1 addresses are leaf addresses — collected in Phase 2

    # Phase 2 — collect all leaf addresses from level-1 nodes
    leaf_addresses = []
    for node in self.all_nodes.get(1, []):
        leaf_addresses.extend(node["addresses"])

    if not leaf_addresses:
        # Depth-0 tree: root is also the leaf, already read in _read_root_node
        return

    # Fetch all leaf nodes in one call
    leaf_size = self._leaf_node_size()
    raw_buffers = self._fetch_fn(leaf_addresses, leaf_size)

    for addr, raw in zip(leaf_addresses, raw_buffers):
        node = self._parse_node_from_buffer(raw, addr, node_level=0)
        self._add_node(node)

3. Leaf node size

HDF5 v1 B-tree nodes carry no explicit size field. The size is a function of
entries_used (from the node header) and the per-entry record size (which depends on
self.dims).

The safest approach is to read the first leaf node's header sequentially to get
entries_used, compute the size, and assume all leaf nodes use the same value. This
is valid because HDF5 fills nodes to capacity before creating a new one; only the
final leaf may have fewer entries. Over-reading that last node is harmless —
_parse_node_from_buffer reads only entries_used entries from the buffer.

def _leaf_node_size(self):
    """
    Compute the byte size of a leaf node by peeking at the first one.

    Layout (NODE_TYPE=1, raw data chunks):
        Header  : 24 bytes  (signature 4 + node_type 1 + node_level 1 +
                              entries_used 2 + left_sibling 8 + right_sibling 8)
        Per entry:
            chunk_size + filter_mask : 8 bytes  (two uint32)
            chunk_offset             : 8 * dims bytes  (dims uint64)
            chunk_address            : 8 bytes  (uint64)
    """
    HEADER_SIZE = 24
    entry_size = 8 + 8 * self.dims + 8

    first_leaf_addr = self.all_nodes[1][0]["addresses"][0]
    self.fh.seek(first_leaf_addr)
    header_bytes = self.fh.read(HEADER_SIZE)
    entries_used = struct.unpack_from("<H", header_bytes, 6)[0]

    return HEADER_SIZE + entries_used * entry_size

4. Parse a node from a raw buffer

Extract the parse logic from _read_node / _read_node_header into a
buffer-based variant. This is also the right place to adopt the memoryview
approach from PR #154 (single bulk read + offset arithmetic) since we already have
the full node bytes in hand:

def _parse_node_from_buffer(self, raw, addr, node_level):
    """
    Parse a BTreeV1RawDataChunks node from a raw bytes buffer.
    Used for parallel leaf fetches; mirrors _read_node / _read_node_header
    but operates on an in-memory buffer rather than seeking self.fh.
    """
    HEADER_SIZE = 24
    # Unpack header
    sig, node_type, nl, entries_used, left_sib, right_sib = struct.unpack_from(
        "<4sBBHQQ", raw, 0
    )
    assert sig == b"TREE"
    assert node_type == self.NODE_TYPE
    assert nl == node_level

    node = {
        "signature":     sig,
        "node_type":     node_type,
        "node_level":    nl,
        "entries_used":  entries_used,
        "left_sibling":  left_sib,
        "right_sibling": right_sib,
    }

    # Unpack entries using memoryview for efficiency (avoids repeated buffer copies)
    mv = memoryview(raw)
    pos = HEADER_SIZE
    entry_size = 8 + 8 * self.dims + 8
    offset_fmt = f"<{self.dims}Q"
    offset_fmt_size = self.dims * 8

    keys = []
    addresses = []
    for _ in range(entries_used):
        chunk_size, filter_mask = struct.unpack_from("<II", mv, pos)
        pos += 8
        chunk_offset = struct.unpack_from(offset_fmt, mv, pos)
        pos += offset_fmt_size
        chunk_address = struct.unpack_from("<Q", mv, pos)[0]
        pos += 8
        keys.append(OrderedDict((
            ("chunk_size",   chunk_size),
            ("filter_mask",  filter_mask),
            ("chunk_offset", chunk_offset),
        )))
        addresses.append(chunk_address)

    node["keys"] = keys
    node["addresses"] = addresses
    self.last_offset = max(addr, self.last_offset)
    return node

---

Changes to DatasetID

DatasetID constructs fetch_fn from its existing parallelism configuration and
passes it when instantiating BTreeV1RawDataChunks. The three cases mirror the
existing dispatch in _select_chunks:

def _make_btree_fetch_fn(self):
    """
    Return a fetch_fn suitable for BTreeV1RawDataChunks, or None for serial fallback.

    fetch_fn(addresses: list[int], size: int) -> list[bytes]
    """
    fh = self._fh
    actual_fh = getattr(fh, "fh", fh)

    # Case A: fsspec cat_ranges
    if self._cat_range_allowed:
        fs  = getattr(actual_fh, "fs",   None)
        path = getattr(actual_fh, "path", None)
        if fs is not None and path is not None and hasattr(fs, "cat_ranges"):
            def fetch_cat_ranges(addresses, size):
                stops = [a + size for a in addresses]
                return fs.cat_ranges([path] * len(addresses), addresses, stops)
            return fetch_cat_ranges

    # Case B: os.pread thread pool
    if self.posix and hasattr(os, "pread") and self._thread_count != 0:
        filename = self._filename
        thread_count = self._thread_count
        def fetch_pread(addresses, size):
            fh_local = open(filename, "rb")
            fd = fh_local.fileno()
            try:
                with ThreadPoolExecutor(max_workers=thread_count) as executor:
                    return list(executor.map(
                        lambda addr: os.pread(fd, size, addr), addresses
                    ))
            finally:
                fh_local.close()
        return fetch_pread

    # Case C: serial — let BTreeV1RawDataChunks handle it internally
    return None

Then wherever BTreeV1RawDataChunks is instantiated in DatasetID, pass the
callable:

# Before:
btree = BTreeV1RawDataChunks(self._fh, offset, dims)

# After:
btree = BTreeV1RawDataChunks(self._fh, offset, dims,
                              fetch_fn=self._make_btree_fetch_fn())

Search DatasetID (and any other callers) for all instantiation sites and update
each one.

---

Interaction with MetadataBufferingWrapper

fetch_fn for the cat_ranges case calls fs.cat_ranges on the underlying fsspec
filesystem, bypassing MetadataBufferingWrapper — exactly as _read_bulk_fsspec
does for chunk data. This is correct:

• Well-packed files: the 1 MiB eager buffer already holds all leaf nodes; the
  fsspec cache layer serves them without network I/O.
• Poorly-packed files: cat_ranges issues concurrent remote fetches for all
  scattered leaf nodes, far better than the current sequential fallback.

No changes to MetadataBufferingWrapper are required.

---

Compatibility and Fallback

• fetch_fn=None is the default. All existing callers that do not pass fetch_fn
  get identical behaviour to today — the inherited sequential _read_children runs
  unchanged.
• The os.pread case gives POSIX users (including pathological local files like
  uas_3hr_IPSL-CM6A-LR_... with 292,192 scattered B-tree nodes) parallel leaf
  fetching without any remote storage requirement.
• All three cases produce identical parsed output; they differ only in how bytes
  arrive.

---

Testing

1. Unit — fetch_fn is called correctly: mock fetch_fn, construct
   BTreeV1RawDataChunks with it, assert it is called once with the correct list of
   leaf addresses and the correct size; assert all_nodes[0] matches a reference
   sequential parse.

2. Unit — fetch_fn=None fallback: construct with no fetch_fn; assert the
   sequential path produces identical all_nodes[0].

3. Unit — _make_btree_fetch_fn dispatch: for each of the three cases (fsspec,
   pread, serial) assert the correct closure type is returned given the parallelism
   configuration.

4. Regression: run the full existing B-tree test suite; all results must be
   bit-identical before and after.

5. Integration — POSIX: open uas_3hr_IPSL-CM6A-LR_... locally with
   _thread_count > 0; assert chunk index matches serial parse; measure wall time.

6. Integration — remote: open the same file via s3fs with cat_range_allowed=True;
   assert chunk index matches serial parse.

---

Files to Change

File	Change
pyfive/btree.py	Add fetch_fn param to BTreeV1RawDataChunks.__init__; override _read_children; add _leaf_node_size and _parse_node_from_buffer
pyfive/dataset_id.py (or wherever DatasetID lives)	Add _make_btree_fetch_fn; update all BTreeV1RawDataChunks(...) instantiation sites
tests/test_btree_parallel.py	New test file covering cases 1–3 above

No changes required to AbstractBTree, BTreeV1, BTreeV1Groups, BTreeV2,
MetadataBufferingWrapper, or ChunkRead.

[bnlawrence]

ok, all of that is implemented in 6b5ba3e, with remarkable results. One of the more pathological CMIP6 files, which had a btree which took 10s to read (on local SSD), now takes 1s (with 5 threads):

Before:

 f = pyfive.File('uas_3hr_IPSL-CM6A-LR_piControl_r1i1p1f1_gr_187001010300-197001010000.nc')

In [3]: time uas = f['uas']
CPU times: user 575 ms, sys: 1.2 s, total: 1.78 s
Wall time: 10.1 s

After

f = pyfive.File('uas_3hr_IPSL-CM6A-LR_piControl_r1i1p1f1_gr_187001010300-197001010000.nc')

In [4]: time uas = f['uas']
CPU times: user 517 ms, sys: 141 ms, total: 658 ms
Wall time: 1.02 s