Fix handling of artificial fields in RNTupleProcessor chains

tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx

@@ -54,14 +54,19 @@ class RNTupleProcessorDS final : public ROOT::RDF::RDataSource {
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Holds useful information about fields added to the RNTupleProcessorDS
    struct RFieldInfo {
-      ROOT::DescriptorId_t fFieldId;
+      ROOT::DescriptorId_t fOnDiskId;
       std::string fFieldName;
       std::string fTypeName;
       std::size_t fNRepetitions;
+      RNTupleProcessorEntry::FieldIndex_t fEntryFieldIdx;
       // Enable `std::vector::emplace_back` for this type
-      RFieldInfo(ROOT::DescriptorId_t fieldId, std::string_view fieldName, std::string_view typeName,
-                 std::size_t nRepetitions)
-         : fFieldId(fieldId), fFieldName(fieldName), fTypeName(typeName), fNRepetitions(nRepetitions)
+      RFieldInfo(ROOT::DescriptorId_t onDiskId, std::string_view fieldName, std::string_view typeName,
+                 std::size_t nRepetitions, Internal::RNTupleProcessorEntry::FieldIndex_t entryFieldIdx)
+         : fOnDiskId(onDiskId),
+           fFieldName(fieldName),
+           fTypeName(typeName),
+           fNRepetitions(nRepetitions),
+           fEntryFieldIdx(entryFieldIdx)
       {
       }
    };

tree/dataframe/src/RNTupleProcessorDS.cxx

@@ -89,7 +89,9 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie
       // The field is a collection or a fixed-size array.
       // We open a new collection scope with fieldID being the inner most collection. E.g. for "event.tracks.hits",
       // fieldInfos would already contain the fieldID of "event.tracks"
-      fieldInfos.emplace_back(field.GetOnDiskId(), field.GetFieldName(), field.GetTypeName(), nRepetitions);
+      const std::string name = (procProvenance.Empty() ? "" : procProvenance.Get() + ".") + std::string(colName);
+      fieldInfos.emplace_back(field.GetOnDiskId(), field.GetFieldName(), field.GetTypeName(), nRepetitions,
+                              *fProcessor->GetEntry().FindFieldIndex(name));
       nFieldInfosAdded++;
    }
 
@@ -127,7 +129,13 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie
 
    // The fieldID could be the root field or the class of fieldId might not be loaded.
    // In these cases, only the inner fields are exposed as RDF columns.
-   auto fieldOrException = ROOT::RFieldBase::Create(field.GetFieldName(), field.GetTypeName());
+   std::string fieldName = field.GetFieldName();
+   auto fieldInfoIt = fieldInfos.begin();
+   while (fieldName == "_0") {
+      fieldName = fieldInfoIt->fFieldName;
+      fieldInfoIt++;
+   }
+   auto fieldOrException = ROOT::RFieldBase::Create(fieldName, field.GetTypeName());
    if (!fieldOrException)
       return;
    auto valueField = fieldOrException.Unwrap();
@@ -160,7 +168,7 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie
       } else {
          cardinalityField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::size_t>>(name);
       }
-      cardinalityField->SetOnDiskId(info.fFieldId);
+      cardinalityField->SetOnDiskId(info.fOnDiskId);
    }
 
    for (auto i = fieldInfos.rbegin(); i != fieldInfos.rend(); ++i) {
@@ -183,7 +191,7 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie
          }
       }
 
-      valueField->SetOnDiskId(fieldInfo.fFieldId);
+      valueField->SetOnDiskId(fieldInfo.fOnDiskId);
 
       // Skip the inner-most collection level to construct the cardinality column
       // It's taken care of by the `if (!fieldInfos.empty())` scope above
@@ -198,29 +206,36 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie
             cardinalityField = std::make_unique<ROOT::RRVecField>(cardinalityFieldName, cardinalityField->Clone("_0"));
          }
 
-         cardinalityField->SetOnDiskId(fieldInfo.fFieldId);
+         cardinalityField->SetOnDiskId(fieldInfo.fOnDiskId);
       }
    }
 
-   if (cardinalityField) {
-      std::string cardinalityFieldName =
-         (procProvenance.Empty() ? "R_rdf_sizeof_" : procProvenance.Get() + "_R_rdf_sizeof_") + std::string(colName);
-      fColumnNames.emplace_back(cardinalityFieldName);
-      fColumnTypes.emplace_back(cardinalityField->GetTypeName());
-      fProcessor->AddFieldToEntry(std::move(cardinalityField), cardinalityFieldName, nullptr,
-                                  Internal::RNTupleProcessorProvenance());
-   }
-
-   fieldInfos.emplace_back(field.GetOnDiskId(), field.GetFieldName(), field.GetTypeName(), nRepetitions);
-   nFieldInfosAdded++;
    std::string canonicalName = (procProvenance.Empty() ? "" : procProvenance.Get() + ".") + std::string(colName);
    fColumnNames.emplace_back(canonicalName);
    fColumnTypes.emplace_back(valueField->GetTypeName());
 
-   // Add nested fields explicitly to the entry, because they may be mapped differently.
-   if (fieldInfos.size() > 1) {
-      fProcessor->AddFieldToEntry(std::move(valueField), canonicalName, nullptr,
-                                  Internal::RNTupleProcessorProvenance());
+   RNTupleProcessorEntry::FieldIndex_t fieldIdx;
+   if (auto optFieldIdx = fProcessor->GetEntry().FindFieldIndex(canonicalName);
+       optFieldIdx && fProcessor->GetEntry().GetField(*optFieldIdx).GetTypeName() == field.GetTypeName()) {
+      fieldIdx = *optFieldIdx;
+   } else {
+      fieldIdx = fProcessor->AddFieldToEntry(std::move(valueField), canonicalName, nullptr,
+                                             Internal::RNTupleProcessorProvenance());
+   }
+
+   if (cardinalityField) {
+      // FIXME(fdegeus): Make cardinality field name work properly with "#colname"
+      std::string cardinalityFieldName =
+         (procProvenance.Empty() ? "R_rdf_sizeof_" : procProvenance.Get() + ".R_rdf_sizeof_") + std::string(colName);
+      fColumnNames.emplace_back(cardinalityFieldName);
+      fColumnTypes.emplace_back(cardinalityField->GetTypeName());
+      auto cardinalityFieldIdx = fProcessor->AddFieldToEntry(std::move(cardinalityField), cardinalityFieldName, nullptr,
+                                                             Internal::RNTupleProcessorProvenance());
+      if (!fieldInfos.empty()) {
+         const auto &fieldInfo = fieldInfos.back();
+         fieldIdx = fieldInfo.fEntryFieldIdx;
+      }
+      fProcessor->GetEntry().SetFieldOriginIdx(cardinalityFieldIdx, fieldIdx);
    }
 
    for (std::uint16_t i = 0; i < nFieldInfosAdded; i++) {

tree/dataframe/test/datasource_ntuple_processor.cxx

@@ -6,6 +6,8 @@
 #include <ROOT/RNTupleWriter.hxx>
 #include <ROOT/RPageStorage.hxx>
 
+#include <TFile.h>
+
 #include <ROOT/TestSupport.hxx>
 
 #include <gtest/gtest.h>
@@ -290,31 +292,39 @@ TEST_F(RNTupleProcessorDSTest, Chain)
    ChainTest(fNtplName, fFileName);
 }
 
-TEST(RNTupleProcessorDS, Join)
-{
-   FileRaii guardFile1("RNTupleProcessorDS_test_join_1.root");
-   FileRaii guardFile2("RNTupleProcessorDS_test_join_2.root");
+class RNTupleProcessorDSAlignedJoinsTest : public ::testing::Test {
+protected:
+   std::vector<std::string> fFiles = {"RNTupleProcessorDSJoins1.root", "RNTupleProcessorDSJoins2.root"};
 
+   void SetUp() override
    {
-      auto model = RNTupleModel::Create();
-      auto fldElectron = model->MakeField<Electron>("electrons");
-      auto writer = RNTupleWriter::Recreate(std::move(model), "primary", guardFile1.GetPath());
-      for (unsigned i = 0; i < 10; ++i) {
-         fldElectron->pt = static_cast<float>(i);
-         writer->Fill();
-      }
-   }
-   {
-      auto model = RNTupleModel::Create();
-      auto fldJets = model->MakeField<std::vector<float>>("jets");
-      auto writer = RNTupleWriter::Recreate(std::move(model), "auxiliary", guardFile2.GetPath());
-      for (unsigned i = 0; i < 10; ++i) {
-         *fldJets = {static_cast<float>(i) * .1f, static_cast<float>(i) * .2f};
-         writer->Fill();
+      for (unsigned i = 0; i < fFiles.size(); i++) {
+         {
+            auto model = RNTupleModel::Create();
+            auto fldElectron = model->MakeField<Electron>("electrons");
+            auto writer = RNTupleWriter::Recreate(std::move(model), "primary", fFiles[i]);
+            for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) {
+               fldElectron->pt = static_cast<float>(n);
+               writer->Fill();
+            }
+         }
+         {
+            auto file = std::unique_ptr<TFile>(TFile::Open(fFiles[i].c_str(), "UPDATE"));
+            auto model = RNTupleModel::Create();
+            auto fldJets = model->MakeField<std::vector<float>>("jets");
+            auto writer = RNTupleWriter::Append(std::move(model), "auxiliary", *file);
+            for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) {
+               *fldJets = {static_cast<float>(n) * .1f, static_cast<float>(n) * .2f};
+               writer->Fill();
+            }
+         }
       }
    }
+};
 
-   auto proc = CreateRNTupleJoinProcessor({"primary", guardFile1.GetPath()}, {"auxiliary", guardFile2.GetPath()}, {});
+TEST_F(RNTupleProcessorDSAlignedJoinsTest, SingleJoin)
+{
+   auto proc = CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {});
    auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc));
 
    auto count = df.Count();
@@ -330,42 +340,102 @@ TEST(RNTupleProcessorDS, Join)
       },
       [](float a, float b) { return a + b; }, "auxiliary.jets");
    EXPECT_FLOAT_EQ(13.5, jetSum.GetValue());
+
+   auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets");
+   EXPECT_FLOAT_EQ(20, nJetSum.GetValue());
 }
 
-TEST(RNTupleProcessorDS, JoinUnaligned)
+TEST_F(RNTupleProcessorDSAlignedJoinsTest, ChainedJoin)
 {
-   FileRaii guardFile1("RNTupleProcessorDS_test_join_unaligned_1.root");
-   FileRaii guardFile2("RNTupleProcessorDS_test_join_unaligned_2.root");
+   std::vector<std::unique_ptr<RNTupleProcessor>> joins;
+   joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {}));
+   joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[1]}, {"auxiliary", fFiles[1]}, {}));
+   auto chain = RNTupleProcessor::CreateChain(std::move(joins));
+   auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(chain));
 
+   auto count = df.Count();
+   EXPECT_EQ(20, count.GetValue());
+
+   auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count();
+   EXPECT_EQ(15, electronFilterCount.GetValue());
+
+   auto jetSum = df.Aggregate(
+      [](float &acc, const ROOT::RVec<float> &jets) {
+         for (const auto &jet : jets)
+            acc += jet;
+      },
+      [](float a, float b) { return a + b; }, "auxiliary.jets");
+   EXPECT_FLOAT_EQ(57, jetSum.GetValue());
+
+   auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets");
+   EXPECT_FLOAT_EQ(40, nJetSum.GetValue());
+}
+
+TEST_F(RNTupleProcessorDSAlignedJoinsTest, JoinedChain)
+{
+   auto chain_primary = RNTupleProcessor::CreateChain({{"primary", fFiles[0]}, {"primary", fFiles[1]}});
+   auto chain_auxiliary = RNTupleProcessor::CreateChain({{"auxiliary", fFiles[0]}, {"auxiliary", fFiles[1]}});
+   auto join = CreateRNTupleJoinProcessor(std::move(chain_primary), std::move(chain_auxiliary), {});
+   auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(join));
+
+   auto count = df.Count();
+   EXPECT_EQ(20, count.GetValue());
+
+   auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count();
+   EXPECT_EQ(15, electronFilterCount.GetValue());
+
+   auto jetSum = df.Aggregate(
+      [](float &acc, const ROOT::RVec<float> &jets) {
+         for (const auto &jet : jets)
+            acc += jet;
+      },
+      [](float a, float b) { return a + b; }, "auxiliary.jets");
+   EXPECT_FLOAT_EQ(57, jetSum.GetValue());
+
+   auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets");
+   EXPECT_FLOAT_EQ(40, nJetSum.GetValue());
+}
+
+class RNTupleProcessorDSUnalignedJoinsTest : public ::testing::Test {
+protected:
+   std::vector<std::string> fFiles = {"RNTupleProcessorDSJoins1.root", "RNTupleProcessorDSJoins2.root"};
+
+   void SetUp() override
    {
-      auto model = RNTupleModel::Create();
-      auto fldId = model->MakeField<std::uint32_t>("idx");
-      auto fldElectron = model->MakeField<Electron>("electrons");
-      auto writer = RNTupleWriter::Recreate(std::move(model), "primary", guardFile1.GetPath());
-      for (unsigned i = 0; i < 10; ++i) {
-         // only fill odd events
-         if (i % 2 == 0)
-            continue;
-         *fldId = i;
-         fldElectron->pt = static_cast<float>(i);
-         writer->Fill();
-      }
-   }
-   {
-      auto model = RNTupleModel::Create();
-      auto fldId = model->MakeField<std::uint32_t>("idx");
-      auto fldJets = model->MakeField<std::vector<float>>("jets");
-      auto writer = RNTupleWriter::Recreate(std::move(model), "auxiliary", guardFile2.GetPath());
-      // fill in reverse
-      for (int i = 9; i >= 0; --i) {
-         *fldId = i;
-         *fldJets = {static_cast<float>(i) * .1f, static_cast<float>(i) * .2f};
-         writer->Fill();
+      for (unsigned i = 0; i < fFiles.size(); i++) {
+         {
+            auto model = RNTupleModel::Create();
+            auto fldIdx = model->MakeField<int>("idx");
+            auto fldElectron = model->MakeField<Electron>("electrons");
+            auto writer = RNTupleWriter::Recreate(std::move(model), "primary", fFiles[i]);
+            for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) {
+               // only fill odd events
+               if (n % 2 == 0)
+                  continue;
+               *fldIdx = n;
+               fldElectron->pt = static_cast<float>(n);
+               writer->Fill();
+            }
+         }
+         {
+            auto file = std::unique_ptr<TFile>(TFile::Open(fFiles[i].c_str(), "UPDATE"));
+            auto model = RNTupleModel::Create();
+            auto fldIdx = model->MakeField<int>("idx");
+            auto fldJets = model->MakeField<std::vector<float>>("jets");
+            auto writer = RNTupleWriter::Append(std::move(model), "auxiliary", *file);
+            for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) {
+               *fldIdx = n;
+               *fldJets = {static_cast<float>(n) * .1f, static_cast<float>(n) * .2f};
+               writer->Fill();
+            }
+         }
       }
    }
+};
 
-   auto proc =
-      CreateRNTupleJoinProcessor({"primary", guardFile1.GetPath()}, {"auxiliary", guardFile2.GetPath()}, {"idx"});
+TEST_F(RNTupleProcessorDSUnalignedJoinsTest, SingleJoin)
+{
+   auto proc = CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {"idx"});
    auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc));
 
    auto count = df.Count();
@@ -381,6 +451,60 @@ TEST(RNTupleProcessorDS, JoinUnaligned)
       },
       [](float a, float b) { return a + b; }, "auxiliary.jets");
    EXPECT_FLOAT_EQ(7.5, jetSum.GetValue());
+
+   auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets");
+   EXPECT_FLOAT_EQ(10, nJetSum.GetValue());
+}
+
+TEST_F(RNTupleProcessorDSUnalignedJoinsTest, ChainedJoin)
+{
+   std::vector<std::unique_ptr<RNTupleProcessor>> joins;
+   joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {"idx"}));
+   joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[1]}, {"auxiliary", fFiles[1]}, {"idx"}));
+   auto chain = RNTupleProcessor::CreateChain(std::move(joins));
+   auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(chain));
+
+   auto count = df.Count();
+   EXPECT_EQ(10, count.GetValue());
+
+   auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count();
+   EXPECT_EQ(8, electronFilterCount.GetValue());
+
+   auto jetSum = df.Aggregate(
+      [](float &acc, const ROOT::RVec<float> &jets) {
+         for (const auto &jet : jets)
+            acc += jet;
+      },
+      [](float a, float b) { return a + b; }, "auxiliary.jets");
+   EXPECT_FLOAT_EQ(30, jetSum.GetValue());
+
+   auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets");
+   EXPECT_FLOAT_EQ(20, nJetSum.GetValue());
+}
+
+TEST_F(RNTupleProcessorDSUnalignedJoinsTest, JoinedChain)
+{
+   auto chain_primary = RNTupleProcessor::CreateChain({{"primary", fFiles[0]}, {"primary", fFiles[1]}});
+   auto chain_auxiliary = RNTupleProcessor::CreateChain({{"auxiliary", fFiles[0]}, {"auxiliary", fFiles[1]}});
+   auto join = CreateRNTupleJoinProcessor(std::move(chain_primary), std::move(chain_auxiliary), {"idx"});
+   auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(join));
+
+   auto count = df.Count();
+   EXPECT_EQ(10, count.GetValue());
+
+   auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count();
+   EXPECT_EQ(8, electronFilterCount.GetValue());
+
+   auto jetSum = df.Aggregate(
+      [](float &acc, const ROOT::RVec<float> &jets) {
+         for (const auto &jet : jets)
+            acc += jet;
+      },
+      [](float a, float b) { return a + b; }, "auxiliary.jets");
+   EXPECT_FLOAT_EQ(30, jetSum.GetValue());
+
+   auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets");
+   EXPECT_FLOAT_EQ(20, nJetSum.GetValue());
 }
 
 #ifdef R__USE_IMT