[TDF] Use SelectColumns instead of PickBranchNames

SelectColumns chooses between the optional user-provided set of
column names and the default set. Throws in case of error. Picks
an incremental number of column names from the default set until
the required number is reached, instead of requiring that the
number of default branches is strictly equal to the number of
required branches for the transformation/action.

tree/treeplayer/inc/ROOT/TDFInterface.hxx

@@ -203,7 +203,7 @@ public:
       auto df = GetDataFrameChecked();
       const ColumnNames_t &defBl = df->GetDefaultColumnNames();
       auto nArgs = TTraits::CallableTraits<F>::arg_types::list_size;
-      const ColumnNames_t &actualBl = TDFInternal::PickBranchNames(nArgs, bn, defBl);
+      const auto actualBl = TDFInternal::SelectColumns(nArgs, bn, defBl);
       using DFF_t = TDFDetail::TFilter<F, Proxied>;
       auto FilterPtr = std::make_shared<DFF_t>(std::move(f), actualBl, *fProxiedPtr, name);
       df->Book(FilterPtr);
@@ -282,7 +282,7 @@ public:
       TDFInternal::CheckTmpBranch(name, df->GetTree());
       const ColumnNames_t &defBl = df->GetDefaultColumnNames();
       auto nArgs = TTraits::CallableTraits<F>::arg_types::list_size;
-      const ColumnNames_t &actualBl = TDFInternal::PickBranchNames(nArgs, bl, defBl);
+      const auto actualBl = TDFInternal::SelectColumns(nArgs, bl, defBl);
       using DFB_t = TDFDetail::TCustomColumn<F, Proxied>;
       const std::string nameInt(name);
       auto BranchPtr = std::make_shared<DFB_t>(nameInt, std::move(expression), actualBl, *fProxiedPtr);
@@ -479,7 +479,7 @@ public:
       auto df = GetDataFrameChecked();
       const ColumnNames_t &defBl = df->GetDefaultColumnNames();
       auto nArgs = TTraits::CallableTraits<F>::arg_types::list_size;
-      const ColumnNames_t &actualBl = TDFInternal::PickBranchNames(nArgs - 1, bl, defBl);
+      const auto actualBl = TDFInternal::SelectColumns(nArgs - 1, bl, defBl);
       using Helper_t = TDFInternal::ForeachSlotHelper<F>;
       using Action_t = TDFInternal::TAction<Helper_t, Proxied>;
       df->Book(std::make_shared<Action_t>(Helper_t(std::move(f)), actualBl, *fProxiedPtr));
@@ -525,12 +525,13 @@ public:
       using arg_types = typename TTraits::CallableTraits<F>::arg_types;
       TDFInternal::CheckReduce(f, arg_types());
       auto df = GetDataFrameChecked();
-      unsigned int nSlots = df->GetNSlots();
-      auto bl = GetBranchNames<T>({columnName}, "reduce branch values");
+      const auto &defBl = df->GetDefaultColumnNames();
+      const ColumnNames_t userColumns = columnName.empty() ? ColumnNames_t() : ColumnNames_t({std::string(columnName)});
+      const auto actualBl = TDFInternal::SelectColumns(1, userColumns, defBl);
       auto redObjPtr = std::make_shared<T>(initValue);
       using Helper_t = TDFInternal::ReduceHelper<F, T>;
       using Action_t = typename TDFInternal::TAction<Helper_t, Proxied>;
-      df->Book(std::make_shared<Action_t>(Helper_t(std::move(f), redObjPtr, nSlots), bl, *fProxiedPtr));
+      df->Book(std::make_shared<Action_t>(Helper_t(std::move(f), redObjPtr, df->GetNSlots()), actualBl, *fProxiedPtr));
       return MakeResultProxy(redObjPtr, df);
    }
 

tree/treeplayer/inc/ROOT/TDFUtils.hxx

@@ -182,7 +182,7 @@ void CheckReduce(F &, T)
 }
 
 /// Returns local BranchNames or default BranchNames according to which one should be used
-const ColumnNames_t &PickBranchNames(unsigned int nArgs, const ColumnNames_t &bl, const ColumnNames_t &defBl);
+const ColumnNames_t SelectColumns(unsigned int nArgs, const ColumnNames_t &bl, const ColumnNames_t &defBl);
 
 namespace ActionTypes {
 struct Histo1D {

tree/treeplayer/src/TDFUtils.cxx

@@ -143,23 +143,25 @@ void CheckTmpBranch(std::string_view branchName, TTree *treePtr)
    }
 }
 
-/// Returns local BranchNames or default BranchNames according to which one should be used
-const ColumnNames_t &PickBranchNames(unsigned int nArgs, const ColumnNames_t &bl, const ColumnNames_t &defBl)
+/// Choose between local column names or default column names, throw in case of errors.
+const ColumnNames_t SelectColumns(unsigned int nRequiredNames, const ColumnNames_t &names,
+                                  const ColumnNames_t &defaultNames)
 {
-   bool useDefBl = false;
-   if (nArgs != bl.size()) {
-      if (bl.size() == 0 && nArgs == defBl.size()) {
-         useDefBl = true;
-      } else {
-         auto msg = "mismatch between number of filter/define arguments (" + std::to_string(nArgs) +
-                    ") and number of columns specified (" + std::to_string(bl.size() ? bl.size() : defBl.size()) +
-                    "). Please check the number of arguments of the function/lambda/functor and the number of branches "
-                    "specified.";
-         throw std::runtime_error(msg);
-      }
+   if (names.empty()) {
+      // use default column names
+      if (defaultNames.size() < nRequiredNames)
+         throw std::runtime_error(std::to_string(nRequiredNames) +
+                                  " column names are required but none were provided and the default list has size " +
+                                  std::to_string(defaultNames.size()));
+      // return first nRequiredNames default column names
+      return ColumnNames_t(defaultNames.begin(), defaultNames.begin() + nRequiredNames);
+   } else {
+      // use column names provided by the user to this particular transformation/action
+      if (names.size() != nRequiredNames)
+         throw std::runtime_error(std::to_string(nRequiredNames) + " column names are required but " +
+                                  std::to_string(names.size()) + " were provided.");
+      return names;
    }
-
-   return useDefBl ? defBl : bl;
 }
 
 } // end NS TDF