diff --git a/src/coffea/nanoevents/mapping/parquet.py b/src/coffea/nanoevents/mapping/parquet.py index 8ba12a1a6..73d8c379f 100644 --- a/src/coffea/nanoevents/mapping/parquet.py +++ b/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): diff --git a/tests/test_nanoevents.py b/tests/test_nanoevents.py index cf48c561e..7b815c87b 100644 --- a/tests/test_nanoevents.py +++ b/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).