Skip to content

Make using DataSource faster by using lazy reading#949

Open
jmcarcell wants to merge 4 commits intomasterfrom
datasource
Open

Make using DataSource faster by using lazy reading#949
jmcarcell wants to merge 4 commits intomasterfrom
datasource

Conversation

@jmcarcell
Copy link
Copy Markdown
Member

@jmcarcell jmcarcell commented Mar 26, 2026

The problem of DataSource is that everything is being loaded all the time, unlike for traditional RDataFrames. I have added lazy reading to ROOTReader and RNTupleReader, which have to be different because the internals of the readers are different. For ROOTReader it is relatively trivial with a callback, we just read the collection that we need when it is used; for RNTupleReader we create a new ROOT RNTupleReader with a minimal model for each collection since the complete model is fixed at the beginning (the key difference is: for ROOTReader we have one branch for each collection and we read per-branch, for RNTupleReader we have the full model). Possible questions and comments:

  • If we want to support lazy reading more generally, for example for files for which we are only interested in a few collections. Then we would maybe need to have different modes for the readers (at least for the RNTuple one this has to be known before reading). Currently, the implementation of the ROOTReader assumes all reading for an event is done before moving on to the next one, which is something that would have to be changed.
  • In that case, maybe lazy reading could be the default if performance is not too different from what we have now.
  • If we don't, then the lazy reading has to be hidden and only allowed for DataSource, since it does not work in the general case (at least for ROOTReader, explained above).

Benchmarks later but I can read a few GB of TTree and RNTuple files in a close time to using RDataFrame directly on them (tested with single threading only, I think multithreading should bring a similar speedup since DataSource makes several independent readers).

BEGINRELEASENOTES

ENDRELEASENOTES

@tmadlener
Copy link
Copy Markdown
Collaborator

When I initially designed the whole Frame infrastructure, I had the following idea for reading data lazily: Construct some form of LazyFrameData that effectively retains a reference to the reader such that it can read the buffers from there when they are requested from the Frame. For ROOT this would effectively mean that the reader does almost nothing and most of the logic for retrieving data from the file would go into this LazyFrameData.

From the looks of your version here callback goes pretty much in that direction, only that the logic still lives in the Reader and that there is a condition of not going to the next entry before all collections have been lazily read. To make this more generic one would have to probably add

  • a mutex (or another synchronization primitive) to lock the reader from the Frame
  • some bookkeeping information to know which entry in the file belongs to a Frame

Purely from not overloading the existing reader with too much functionality, I would be in favor of having the details of lazy or eager reading entirely hidden behind the existing interface of Reader, i.e. no addition of readFrameLazy, but rather have new lazy readers (and corresponding lazy frames) that provide the same interface.


As a side note some of the excessive data loading could be front-loaded to the users for a quick work around, because DataSource has the ability to only read a subset of all collections. @kjvbrt and me had a brief discussion about automating that on the python side, but I think we arrived at the conclusion that that would essentially require parsing python or doing a "dummy" run to collect data names. However, a user might know which collections they want and could provide that as a list.

@jmcarcell
Copy link
Copy Markdown
Member Author

Making a LazyFrameData and new readers seems like the "lazy" option. Currently the callback can be inserted with few changes in the ROOTFrameData, that is the same except for the callback, and the runtime performance is simply an if check when not using it. The readers are the same except that they need the callback (and the RNTuple to create a map or list of "small" readers). Even the readLazy functions are almost a copy and paste of the existing readEntry functions. So making new readers duplicates all the code of the existing readers. Another way to go at this could be to pass a parameter "lazy" to the readers at construction time and then enable the lazy functionality and have two paths that share a lot of the functionality. If lazy -> init the lazy path (only different for the RNTupleReader) and read lazily instead of the regular way (different for both readers since they have to use the callback).

I think performance for the ROOTReader could be the same (note that nothing changes for initialization unlike the RNTupleReader that freezes the model) as the one that we have, since it is reading collections one by one. So it wouldn't be crazy if lazy reading was the default.

@tmadlener
Copy link
Copy Markdown
Collaborator

The thing I don't like about the current implementation of readLazy is that it couples the returned FrameData to the state of the reader. It's even explicitly mentioned in the docstring:

/// The reader must not advance to the next entry while the returned FrameData
/// (and any Frame built from it) is still being accessed. This is guaranteed
/// in the DataSource context where each slot has its own reader.

If we want to make this generally useful, what I would like to guarantee is that the following is also possible and works as expected

auto reader = podio::makeReader("some-file.root"); // <-- Assume for now this is a lazy reader through magic

auto event1 = reader.readEvent(1);
auto event2 = reader.readNextEvent();
auto event3 = reader.readEvent(42);

auto coll1 = event1.get<CollType>("collName");
auto coll2 = event2.get<AnotherType>("name");
auto coll3 = event3.get<ThirdType>("third-name");

In the approach in this PR, at the moment this would not work and presumably break somewhere. Instead if there is a dedicated LazyFrameData (or probably one per backend, so two for ROOT) that just keeps track of the original reader that created it as well as which entry in that reader this can be made to work. In the end for the users it should not make any difference in usage whether they use lazy or eager reading and this should be drop-in replaceable. Otherwise, I don't think we have gained much by adding lazy reading.

Comment on lines +181 to +200
// Build a minimal RNTupleModel from the full reader's descriptor
auto& fullReader = *m_readers[category][readerIndex];
const auto& desc = fullReader.GetDescriptor();

ROOT::RCreateFieldOptions fieldOpts;
fieldOpts.SetEmulateUnknownTypes(true);
fieldOpts.SetReturnInvalidOnError(true);

auto smallModel = ROOT::RNTupleModel::CreateBare();
const auto& topFieldDesc = desc.GetFieldDescriptor(desc.GetFieldZeroId());
for (const auto& fieldDesc : desc.GetFieldIterable(topFieldDesc)) {
const auto& fn = fieldDesc.GetFieldName();
if (std::ranges::find(neededFieldNames, fn) != neededFieldNames.end()) {
auto field = fieldDesc.CreateField(desc, fieldOpts);
if (field) {
smallModel->AddField(std::move(field));
}
}
}
smallModel->Freeze();
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we had a "global" (per reader) lazy flag this could move into initCategory, right? In that case one could simply create all readers up-front in initCategory or even up-front in openFiles like we do with the full reader and simply create one per possible collection?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants