[TDF] Define DataSource columns on demand in CreateActions

Most actions (the ones that are handled by CreateActions) can now
detect when a column they act upon comes from a datasource and has
not yet been defined. The appropriate `Define`ition is triggered
accordingly.

Author: eguiraud

tree/treeplayer/inc/ROOT/TDFInterface.hxx

@@ -1149,13 +1149,15 @@ private:
              typename std::enable_if<!TDFInternal::TNeedJitting<BranchTypes...>::value, int>::type = 0>
    TResultProxy<ActionResultType> CreateAction(const ColumnNames_t &columns, const std::shared_ptr<ActionResultType> &r)
    {
-      auto loopManager = GetDataFrameChecked();
-      const auto nColumns = sizeof...(BranchTypes);
-      const auto validColumnNames = GetValidatedColumnNames(*loopManager, nColumns, columns);
+      auto lm = GetDataFrameChecked();
+      constexpr auto nColumns = sizeof...(BranchTypes);
+      const auto colsAndUndefCols = CheckColumnNames(columns, nColumns, *lm);
+      if (fDataSource)
+         DefineDataSourceColumns<BranchTypes...>(colsAndUndefCols, *lm, TDFInternal::GenStaticSeq_t<nColumns>());
+      const auto &theColumnNames = colsAndUndefCols.first;
       const auto nSlots = fProxiedPtr->GetNSlots();
-      TDFInternal::BuildAndBook<BranchTypes...>(validColumnNames, r, nSlots, *loopManager, *fProxiedPtr,
-                                                (ActionType *)nullptr);
-      return MakeResultProxy(r, loopManager);
+      TDFInternal::BuildAndBook<BranchTypes...>(theColumnNames, r, nSlots, *lm, *fProxiedPtr, (ActionType *)nullptr);
+      return MakeResultProxy(r, lm);
    }
 
    // User did not specify type, do type inference
@@ -1238,12 +1240,18 @@ private:
       return snapshotTDF;
    }
 
-   ColumnNames_t GetValidatedColumnNames(TLoopManager &lm, const unsigned int nColumns,
-                                         const ColumnNames_t &userColumns)
+   /// Given the desired number of columns and the user-provided list of columns:
+   /// * fallback to using the first nColumns default columns if needed (or throw if nColumns > nDefaultColumns)
+   /// * check that selected column names refer to valid branches, custom columns or datasource columns (throw if not)
+   /// Return the list of selected column names and a bitmask indicating the columns that must be defined via datasource
+   std::pair<ColumnNames_t, std::vector<bool>>
+   CheckColumnNames(const ColumnNames_t &userColumns, const unsigned int nColumns, TLoopManager &lm)
    {
       const auto &defaultColumns = lm.GetDefaultColumnNames();
-      const auto trueColumns = TDFInternal::SelectColumns(nColumns, userColumns, defaultColumns);
-      const auto unknownColumns = TDFInternal::FindUnknownColumns(trueColumns, lm.GetTree(), fValidCustomColumns);
+      const auto selectedColumns = TDFInternal::SelectColumns(nColumns, userColumns, defaultColumns);
+      const auto unknownColumns =
+         TDFInternal::FindUnknownColumns(selectedColumns, lm.GetTree(), fValidCustomColumns,
+                                         fDataSource ? fDataSource->GetColumnNames() : ColumnNames_t{});
 
       if (!unknownColumns.empty()) {
          // throw
@@ -1256,7 +1264,54 @@ private:
          throw std::runtime_error("Unknown column" + unknowns.str());
       }
 
-      return trueColumns;
+      const auto mustBeDefined = FindUndefinedDSColumns(selectedColumns, lm.GetDefinedDataSourceColumns());
+      return std::make_pair(selectedColumns, mustBeDefined);
+   }
+
+   //TODO remove usages of this everywhere
+   ColumnNames_t
+   GetValidatedColumnNames(TLoopManager &lm, const unsigned int nColumns, const ColumnNames_t &userColumns) {
+      return CheckColumnNames(userColumns, nColumns, lm).first;
+   }
+
+      /// Return a bitset each element of which indicates whether the corresponding element in `selectedColumns` is the
+      /// name of a column that must be defined via datasource. All elements of the returned vector are false if no
+      /// data-source is present.
+      std::vector<bool> FindUndefinedDSColumns(const ColumnNames_t &requestedCols, const ColumnNames_t &definedDSCols)
+   {
+      const auto nColumns = requestedCols.size();
+      std::vector<bool> mustBeDefined(nColumns, false);
+      if (fDataSource) {
+         for (auto i = 0u; i < nColumns; ++i)
+            mustBeDefined[i] =
+               std::find(definedDSCols.begin(), definedDSCols.end(), requestedCols[i]) == definedDSCols.end();
+      }
+      return mustBeDefined;
+   }
+
+   template <typename... ColumnTypes, int... S>
+   void DefineDataSourceColumns(const std::pair<std::vector<std::string>, std::vector<bool>> &colsAndUndefCols,
+                                TLoopManager &lm, TDFInternal::StaticSeq<S...> /*dummy*/)
+   {
+      assert(fDataSource != nullptr);
+      const auto &columns = colsAndUndefCols.first;
+      const auto &mustBeDefined = colsAndUndefCols.second;
+      if (std::none_of(mustBeDefined.begin(), mustBeDefined.end(), [](bool b) { return b; })) {
+         // no need to define any column
+         return;
+      } else {
+         // hack to expand a template parameter pack without c++17 fold expressions.
+         std::initializer_list<int> expander{
+            (mustBeDefined[S] ? DefineDSColumnHelper<ColumnTypes>(columns[S], lm) : /*no-op*/((void)0), 0)...};
+      }
+   }
+
+   template<typename T>
+   void DefineDSColumnHelper(std::string_view name, TLoopManager &lm) {
+      const auto nSlots = fProxiedPtr->GetNSlots();
+      auto readers = fDataSource->GetColumnReaders<T>(name, nSlots);
+      DefineSlot(name, [readers](unsigned int slot) { return **readers[slot]; });
+      lm.AddDataSourceColumn(name);
    }
 
 protected:

tree/treeplayer/inc/ROOT/TDFNodes.hxx

@@ -102,7 +102,8 @@ class TLoopManager : public std::enable_shared_from_this<TLoopManager> {
    unsigned int fNStopsReceived{0}; ///< Number of times that a children node signaled to stop processing entries.
    const ELoopType fLoopType; ///< The kind of event loop that is going to be run (e.g. on ROOT files, on no files)
    std::string fToJit;        ///< string containing all `BuildAndBook` actions that should be jitted before running
-   std::unique_ptr<TDataSource> fDataSource; ///< Owning pointer to a data-source object. Null if no data-source.
+   std::unique_ptr<TDataSource> fDataSource; ///< Owning pointer to a data-source object. Null if no data-source
+   ColumnNames_t fDefinedDataSourceColumns; ///< List of data-source columns that have been `Define`d so far
 
    void RunEmptySourceMT();
    void RunEmptySource();
@@ -149,6 +150,8 @@ public:
    void IncrChildrenCount() { ++fNChildren; }
    void StopProcessing() { ++fNStopsReceived; }
    void Jit(const std::string &s) { fToJit.append(s); }
+   const ColumnNames_t &GetDefinedDataSourceColumns() const { return fDefinedDataSourceColumns; }
+   void AddDataSourceColumn(std::string_view name) { fDefinedDataSourceColumns.emplace_back(name); }
 };
 } // end ns TDF
 } // end ns Detail

tree/treeplayer/inc/ROOT/TDFUtils.hxx

@@ -210,7 +210,8 @@ void CheckSnapshot(unsigned int nTemplateParams, unsigned int nColumnNames);
 const ColumnNames_t SelectColumns(unsigned int nArgs, const ColumnNames_t &bl, const ColumnNames_t &defBl);
 
 /// Check whether column names refer to a valid branch of a TTree or have been `Define`d. Return invalid column names.
-ColumnNames_t FindUnknownColumns(const ColumnNames_t &requiredCols, TTree *tree, const ColumnNames_t &definedCols);
+ColumnNames_t FindUnknownColumns(const ColumnNames_t &requiredCols, TTree *tree, const ColumnNames_t &definedCols,
+                                 const ColumnNames_t &dataSourceColumns);
 
 namespace ActionTypes {
 struct Histo1D {

tree/treeplayer/src/TDFUtils.cxx

@@ -185,7 +185,8 @@ const ColumnNames_t SelectColumns(unsigned int nRequiredNames, const ColumnNames
    }
 }
 
-ColumnNames_t FindUnknownColumns(const ColumnNames_t &requiredCols, TTree *tree, const ColumnNames_t &definedCols)
+ColumnNames_t FindUnknownColumns(const ColumnNames_t &requiredCols, TTree *tree, const ColumnNames_t &definedCols,
+                                 const ColumnNames_t &dataSourceColumns)
 {
    ColumnNames_t unknownColumns;
    for (auto &column : requiredCols) {
@@ -195,6 +196,10 @@ ColumnNames_t FindUnknownColumns(const ColumnNames_t &requiredCols, TTree *tree,
       const auto isCustomColumn = std::find(definedCols.begin(), definedCols.end(), column) != definedCols.end();
       if (isCustomColumn)
          continue;
+      const auto isDataSourceColumn =
+         std::find(dataSourceColumns.begin(), dataSourceColumns.end(), column) != dataSourceColumns.end();
+      if (isDataSourceColumn)
+         continue;
       unknownColumns.emplace_back(column);
    }
    return unknownColumns;