fix(nanoevents): correct parquet entry_start jagged reads and cache column reads

src/coffea/nanoevents/mapping/parquet.py

@@ -3,11 +3,19 @@
 
 import awkward
 import numpy
+from cachetools import LRUCache
 from fsspec.core import OpenFile
 
 from coffea.nanoevents.mapping.base import BaseSourceMapping, UUIDOpener
 from coffea.nanoevents.util import quote, tuple_to_key
 
+# Number of distinct parquet columns kept materialized per open file. A single
+# jagged buffer access (offsets + content) reads the same column more than once,
+# and a NanoEvents view typically touches only a handful of columns at a time,
+# so a small cache eliminates redundant full-column reads without holding the
+# whole file in memory.
+_PARQUET_COLUMN_CACHE_SIZE = 16
+
 
 # IMPORTANT -> For now the uuid is just the uuid of the pfn.
 #              Later we should use the ParquetFile common_metadata to populate.
@@ -25,6 +33,10 @@ def __init__(self, file, dataset=None, openfile: OpenFile = None):
             self.file = file
             self.dataset = dataset
             self.openfile = openfile
+            # Cache materialized single-column tables so that the multiple buffer
+            # accesses for one column (e.g. offsets and content of a jagged array)
+            # do not each re-read the entire column from the parquet file.
+            self._column_cache = LRUCache(_PARQUET_COLUMN_CACHE_SIZE)
 
         def __del__(self):
             """
@@ -36,11 +48,17 @@ def __del__(self):
                 self.openfile.close()
 
         def read(self, column_name):
+            try:
+                return self._column_cache[column_name]
+            except KeyError:
+                pass
             # make sure uproot is single-core since our calling context might not be
             if self.dataset is not None:
-                return self.dataset.to_table(use_threads=False, columns=[column_name])
+                table = self.dataset.to_table(use_threads=False, columns=[column_name])
             else:
-                return self.file.read([column_name], use_threads=False)
+                table = self.file.read([column_name], use_threads=False)
+            self._column_cache[column_name] = table
+            return table
 
         # for right now spoof the notion of directories in files
         # parquet can do it but we've gotta convince people to
@@ -113,16 +131,21 @@ def array(self, entry_start, entry_stop):
             out = None
             if isinstance(aspa, (pa.lib.ListArray, pa.lib.LargeListArray)):
                 value_type = aspa.type.value_type
-                offsets = None
-                if isinstance(aspa, pa.lib.LargeListArray):
-                    offsets = numpy.frombuffer(aspa.buffers()[1], dtype=numpy.int64)[
-                        : len(aspa) + 1
-                    ]
-                else:
-                    offsets = numpy.frombuffer(aspa.buffers()[1], dtype=numpy.int32)[
-                        : len(aspa) + 1
-                    ]
-                    offsets = offsets.astype(numpy.int64)
+                # A sliced pyarrow (Large)ListArray does not copy its buffers; it
+                # only records a logical ``aspa.offset`` into the shared offsets
+                # buffer. Read ``len(aspa) + 1`` offsets starting at that logical
+                # offset and rebase them to start at 0, so the returned
+                # ListOffsetArray indexes into the (already sliced) flattened
+                # content that ``flatten()`` returns.
+                dtype = (
+                    numpy.int64
+                    if isinstance(aspa, pa.lib.LargeListArray)
+                    else numpy.int32
+                )
+                raw_offsets = numpy.frombuffer(aspa.buffers()[1], dtype=dtype)
+                offsets = raw_offsets[aspa.offset : aspa.offset + len(aspa) + 1]
+                offsets = offsets.astype(numpy.int64)
+                offsets = offsets - offsets[0]
                 offsets = awkward.index.Index64(offsets)
 
                 if not isinstance(value_type, pa.lib.DataType):

tests/test_nanoevents.py

@@ -1,4 +1,5 @@
 import os
+from functools import partial
 from pathlib import Path
 
 import awkward as ak
@@ -343,6 +344,160 @@ def test_uproot_write(tmp_path):
     assert ak.all(orig_base.MET_pt == test_base.MET_pt)
 
 
+parquet_suffixes = [
+    "parquet",
+    "extensionarray.parquet",
+]
+
+
+# virtual is the mode the production Runner parquet path uses (executor.py).
+@pytest.mark.parametrize("mode", ["eager", "virtual"])
+@pytest.mark.parametrize("suffix", parquet_suffixes)
+@pytest.mark.parametrize(
+    "entry_start,entry_stop", [(5, 15), (1, 40), (0, 10), (37, 40)]
+)
+def test_parquet_entry_range_matches_full_slice(
+    tests_directory, mode, suffix, entry_start, entry_stop
+):
+    """Reading a parquet file with entry_start > 0 returns the same per-event
+    data as reading the whole file and slicing, for both jagged and flat
+    branches.
+    """
+    path = f"{tests_directory}/samples/nano_dy.{suffix}"
+    from_parquet = getattr(
+        NanoEventsFactory, f"from_{suffix.removeprefix('extensionarray.')}"
+    )
+
+    full = from_parquet(path, schemaclass=NanoAODSchema, mode="eager").events()
+    sub = from_parquet(
+        path,
+        schemaclass=NanoAODSchema,
+        mode=mode,
+        entry_start=entry_start,
+        entry_stop=entry_stop,
+    ).events()
+
+    assert len(sub) == entry_stop - entry_start
+
+    # Jagged collections (the buggy path) must match the full-read slice exactly.
+    for field in ("Muon", "Jet", "Electron"):
+        sub_pt = ak.to_list(getattr(sub, field).pt)
+        full_pt = ak.to_list(getattr(full, field).pt[entry_start:entry_stop])
+        assert (
+            sub_pt == full_pt
+        ), f"{field}.pt mismatch for [{entry_start}:{entry_stop}]"
+
+    # A flat (per-event) branch should match as well.
+    assert ak.to_list(sub.MET.pt) == ak.to_list(full.MET.pt[entry_start:entry_stop])
+
+
+@pytest.mark.parametrize(
+    "entry_start,entry_stop", [(5, 15), (1, 40), (0, 10), (37, 40)]
+)
+def test_parquet_int32_list_offsets_entry_range(tmp_path, entry_start, entry_stop):
+    """Cover the numpy.int32 offsets branch of the parquet entry-range slice.
+
+    The nano_dy sample files all decode to LargeListArray (int64 offsets), so a
+    plain pyarrow ``list_`` column (int32 offsets) is needed to exercise the
+    other side of the dtype selection in ParquetSourceMapping.
+    """
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    n = 40
+    jagged = [[float(i)] * (i % 3) for i in range(n)]
+    table = pa.table(
+        {
+            "jag": pa.array(jagged, type=pa.list_(pa.float32())),
+            "flat": pa.array(np.arange(n, dtype=np.float32)),
+        }
+    )
+    # guard the premise: a non-large list is what yields int32 offsets
+    assert pa.types.is_list(table.schema.field("jag").type)
+    path = str(tmp_path / "int32list.parquet")
+    pq.write_table(table, path)
+
+    full = NanoEventsFactory.from_parquet(
+        path, schemaclass=BaseSchema, mode="eager"
+    ).events()
+    sub = NanoEventsFactory.from_parquet(
+        path,
+        schemaclass=BaseSchema,
+        mode="eager",
+        entry_start=entry_start,
+        entry_stop=entry_stop,
+    ).events()
+
+    assert len(sub) == entry_stop - entry_start
+    assert ak.to_list(sub.jag) == ak.to_list(full.jag[entry_start:entry_stop])
+    assert ak.to_list(sub.flat) == ak.to_list(full.flat[entry_start:entry_stop])
+
+
+def test_parquet_column_cache_avoids_repeated_reads(tests_directory, monkeypatch):
+    """A jagged parquet column is materialized through two separate buffer keys
+    (offsets and content). The per-source column cache collapses these into a
+    single read of the column while returning identical data.
+    """
+    import pyarrow.parquet as pq
+
+    from coffea.nanoevents.factory import _key_formatter
+    from coffea.nanoevents.mapping.parquet import (
+        ParquetSourceMapping,
+        TrivialParquetOpener,
+    )
+    from coffea.nanoevents.util import tuple_to_key
+
+    path = f"{tests_directory}/samples/nano_dy.parquet"
+
+    read_counts = {}
+    orig_read = pq.ParquetFile.read
+
+    def counting_read(self, columns=None, use_threads=True, **kwargs):
+        for c in columns or []:
+            read_counts[c] = read_counts.get(c, 0) + 1
+        return orig_read(self, columns=columns, use_threads=use_threads, **kwargs)
+
+    monkeypatch.setattr(pq.ParquetFile, "read", counting_read)
+
+    parfile = pq.ParquetFile(path)
+    n = parfile.metadata.num_rows
+    mapping = ParquetSourceMapping(TrivialParquetOpener({"uu": path}), 0, n)
+    partition_key = ("uu", "obj", f"0-{n}")
+    mapping.preload_column_source(
+        partition_key[0],
+        partition_key[1],
+        TrivialParquetOpener.UprootLikeShim(parfile),
+    )
+
+    subform = mapping._extract_base_form(parfile.schema_arrow)
+    idx = subform["fields"].index("Muon_pt")
+    jagged_form = {
+        "class": "RecordArray",
+        "fields": ["Muon_pt"],
+        "contents": [subform["contents"][idx]],
+        "parameters": {"__doc__": "parquetfile"},
+        "form_key": "",
+    }
+
+    array = ak.from_buffers(
+        form=ak.forms.from_dict(jagged_form),
+        length=n,
+        container=mapping,
+        buffer_key=partial(_key_formatter, tuple_to_key(partition_key)),
+        highlevel=True,
+    )
+
+    # Offsets and content of one jagged column -> a single underlying read.
+    assert read_counts["Muon_pt"] == 1
+
+    # And the cache must not corrupt the returned data: compare against the
+    # value seen through the normal (un-monkeypatched) reader path.
+    reference = NanoEventsFactory.from_parquet(
+        path, schemaclass=NanoAODSchema, mode="eager"
+    ).events()
+    assert ak.to_list(array.Muon_pt) == ak.to_list(reference.Muon.pt)
+
+
 def test_keys_for_buffer_keys_loadallowmissing():
     """Regression test for scikit-hep/coffea#1578 (bug 5).