Make using DataSource faster by using lazy reading

include/podio/DataSource.h

@@ -7,13 +7,15 @@
 #include <podio/Reader.h>
 
 // ROOT
+#include <ROOT/RDF/RColumnReaderBase.hxx>
 #include <ROOT/RDataFrame.hxx>
 #include <ROOT/RDataSource.hxx>
 
 // STL
 #include <memory>
 #include <string>
 #include <typeinfo>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -100,47 +102,59 @@ class DataSource : public ROOT::RDF::RDataSource {
 
   std::string GetLabel() override {
     return "PODIO Datasource";
-  };
+  }
+
+  // Legacy API
+  std::vector<void*> GetColumnReadersImpl(std::string_view, const std::type_info&) override {
+    return {};
+  }
+
+  std::size_t GetNFiles() const override {
+    return m_filePathList.size();
+  }
 
-protected:
   ///
-  /// @brief Type-erased vector of pointers to pointers to column
-  ///        values --- one per slot.
+  /// @brief Returns a column reader for the given slot and column.
   ///
-  std::vector<void*> GetColumnReadersImpl(std::string_view name, const std::type_info& typeInfo) override;
+  std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase> GetColumnReaders(unsigned int slot, std::string_view name,
+                                                                         const std::type_info& tid) override;
 
+protected:
   std::string AsString() override {
     return "Podio data source";
   }
 
 private:
-  /// Number of slots/threads
-  unsigned int m_nSlots = 1;
-
   /// Input filename
   std::vector<std::string> m_filePathList = {};
 
   /// Total number of events
   ULong64_t m_nEvents = 0;
 
-  /// Ranges of events available to be processed
-  std::vector<std::pair<ULong64_t, ULong64_t>> m_rangesAvailable = {};
-
-  /// Ranges of events available ever created
+  /// All entry ranges, fixed after SetNSlots
   std::vector<std::pair<ULong64_t, ULong64_t>> m_rangesAll = {};
 
+  /// Cursor into m_rangesAll for GetEntryRanges, reset each Initialize()
+  size_t m_rangesCursor = 0;
+
   /// Column names
   std::vector<std::string> m_columnNames{};
 
   /// Column types
   std::vector<std::string> m_columnTypes = {};
 
+  /// Fast column name -> index lookup
+  std::unordered_map<std::string, size_t> m_columnIndex{};
+
   /// Collections, m_Collections[columnIndex][slotIndex]
   std::vector<std::vector<const podio::CollectionBase*>> m_Collections = {};
 
   /// Active collections
   std::vector<unsigned int> m_activeCollections = {};
 
+  /// Names of active collections, kept in sync with m_activeCollections
+  std::vector<std::string> m_activeCollectionNames{};
+
   /// Root podio readers
   std::vector<std::unique_ptr<podio::Reader>> m_podioReaders = {};
 

include/podio/RNTupleReader.h

@@ -6,8 +6,10 @@
 #include "podio/utilities/DatamodelRegistryIOHelpers.h"
 #include "podio/utilities/RootHelpers.h"
 
+#include <map>
 #include <string>
 #include <string_view>
+#include <tuple>
 #include <unordered_map>
 #include <vector>
 
@@ -91,6 +93,15 @@ class RNTupleReader {
   std::unique_ptr<podio::ROOTFrameData> readEntry(std::string_view name, const unsigned entry,
                                                   const std::vector<std::string>& collsToRead = {});
 
+  /// Like readEntry, but collections not in collsToRead are loaded lazily on
+  /// first access rather than eagerly. Also skips reading GenericParameters.
+  ///
+  /// 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.
+  std::unique_ptr<podio::ROOTFrameData> readEntryLazy(const std::string& name, unsigned entry,
+                                                      const std::vector<std::string>& collsToRead);
+
   /// Get the names of all the available Frame categories in the current file(s).
   ///
   /// @returns The names of the available categores from the file
@@ -158,6 +169,9 @@ class RNTupleReader {
   std::unordered_map<std::string_view, std::vector<std::unique_ptr<root_compat::RNTupleReader>>> m_readers{};
   std::unordered_map<std::string, std::unique_ptr<root_compat::RNTupleReader>> m_metadata_readers{};
   std::vector<std::string> m_filenames{};
+  /// Per-category list of filenames (parallel to m_readers[category]); may differ from m_filenames
+  /// if some files don't contain a given category.
+  std::unordered_map<std::string, std::vector<std::string>> m_categoryFilenames{};
 
   std::unordered_map<std::string_view, unsigned> m_entries{};
   // Map category to a vector that contains at how many entries each reader starts
@@ -173,6 +187,17 @@ class RNTupleReader {
   std::vector<std::string> m_availableCategories{};
 
   std::unordered_map<std::string_view, std::shared_ptr<podio::CollectionIDTable>> m_idTables{};
+
+  /// Cache of partial readers, keyed by (category + "|" + sorted-collsToRead-key) + readerIndex.
+  /// Each partial reader is opened with a minimal RNTupleModel containing only the fields
+  /// needed for the requested collections
+  std::map<std::tuple<std::string, std::string, size_t>, std::unique_ptr<root_compat::RNTupleReader>>
+      m_partialReaders{};
+
+  /// Return or lazily create a partial ROOT::RNTupleReader for the given category,
+  /// readerIndex, and set of collections.
+  root_compat::RNTupleReader& getOrCreatePartialReader(const std::string& category, size_t readerIndex,
+                                                       const std::vector<std::string>& collsToRead);
 };
 
 } // namespace podio

include/podio/ROOTFrameData.h

@@ -5,6 +5,7 @@
 #include "podio/CollectionIDTable.h"
 #include "podio/GenericParameters.h"
 
+#include <functional>
 #include <memory>
 #include <optional>
 #include <string>
@@ -17,6 +18,8 @@ class ROOTFrameData {
 
 public:
   using BufferMap = std::unordered_map<std::string, podio::CollectionReadBuffers>;
+  /// Callback type for lazily reading collection buffers on demand
+  using LazyReadFn = std::function<std::optional<podio::CollectionReadBuffers>(const std::string&)>;
 
   ROOTFrameData() = delete;
   ~ROOTFrameData() = default;
@@ -27,6 +30,9 @@ class ROOTFrameData {
 
   ROOTFrameData(BufferMap&& buffers, CollIDPtr&& idTable, podio::GenericParameters&& params);
 
+  /// Constructor with an optional lazy callback for loading collections on-demand
+  ROOTFrameData(BufferMap&& buffers, CollIDPtr&& idTable, podio::GenericParameters&& params, LazyReadFn&& lazyRead);
+
   std::optional<podio::CollectionReadBuffers> getCollectionBuffers(const std::string& name);
 
   podio::CollectionIDTable getIDTable() const;
@@ -42,6 +48,7 @@ class ROOTFrameData {
   // This is co-owned by each FrameData and the original reader. (for now at least)
   CollIDPtr m_idTable{nullptr};
   podio::GenericParameters m_parameters{};
+  LazyReadFn m_lazyRead{nullptr};
 };
 
 } // namespace podio

include/podio/ROOTReader.h

@@ -105,6 +105,15 @@ class ROOTReader {
   std::unique_ptr<podio::ROOTFrameData> readEntry(std::string_view name, const unsigned entry,
                                                   const std::vector<std::string>& collsToRead = {});
 
+  /// Like readEntry, but collections not in collsToRead are loaded lazily on
+  /// first access rather than eagerly. Also skips reading GenericParameters.
+  ///
+  /// 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.
+  std::unique_ptr<podio::ROOTFrameData> readEntryLazy(const std::string& name, unsigned entry,
+                                                      const std::vector<std::string>& collsToRead);
+
   /// Get the number of entries for the given name
   ///
   /// @param name The name of the category

include/podio/Reader.h

@@ -145,6 +145,11 @@ class Reader {
     return m_self->readFrame(name, index, collsToRead);
   }
 
+  /// Like readFrame but loads collections not in collsToRead lazily on first access.
+  /// Falls back to readFrame for non-ROOT backends.
+  /// See ROOTReader::readEntryLazy
+  podio::Frame readFrameLazy(const std::string& name, size_t index, const std::vector<std::string>& collsToRead);
+
   /// Read a specific frame of the "events" category
   ///
   /// @param index The event number to read