[ntuple] Add composed chain tests

root-project · Jan 28, 2025 · fdceaaf · fdceaaf
1 parent 8494e63
commit fdceaaf
Showing 1 changed file with 78 additions and 14 deletions.
diff --git a/tree/ntuple/v7/test/ntuple_processor.cxx b/tree/ntuple/v7/test/ntuple_processor.cxx
@@ -24,27 +24,43 @@ TEST(RNTupleProcessor, EmptyNTuple)
 
 class RNTupleProcessorTest : public testing::Test {
 protected:
-   const std::string fFileName = "test_ntuple_processor.root";
-   const std::string fNTupleName = "ntuple";
+   const std::array<std::string, 2> fFileNames{"test_ntuple_processor1.root", "test_ntuple_processor2.root"};
+   const std::array<std::string, 2> fNTupleNames{"ntuple", "ntuple_aux"};
 
    void SetUp() override
    {
-      auto model = RNTupleModel::Create();
-      auto fldX = model->MakeField<float>("x");
-      auto fldY = model->MakeField<std::vector<float>>("y");
-      auto ntuple = RNTupleWriter::Recreate(std::move(model), fNTupleName, fFileName);
-
-      for (unsigned i = 0; i < 5; i++) {
-         *fldX = static_cast<float>(i);
-         *fldY = {static_cast<float>(i), static_cast<float>(i * 2)};
-         ntuple->Fill();
+      {
+         auto model = RNTupleModel::Create();
+         auto fldI = model->MakeField<int>("i");
+         auto fldX = model->MakeField<float>("x");
+         auto fldY = model->MakeField<std::vector<float>>("y");
+         auto ntuple = RNTupleWriter::Recreate(std::move(model), fNTupleNames[0], fFileNames[0]);
+
+         for (unsigned i = 0; i < 5; i++) {
+            *fldI = i;
+            *fldX = static_cast<float>(i);
+            *fldY = {static_cast<float>(i), static_cast<float>(i * 2)};
+            ntuple->Fill();
+         }
+      }
+      {
+         auto model = RNTupleModel::Create();
+         auto fldI = model->MakeField<int>("i");
+         auto fldZ = model->MakeField<float>("z");
+         auto ntuple = RNTupleWriter::Recreate(std::move(model), fNTupleNames[1], fFileNames[1]);
+
+         for (unsigned i = 0; i < 5; ++i) {
+            *fldI = i;
+            *fldZ = i * 2.f;
+            ntuple->Fill();
+         }
       }
    }
 };
 
 TEST_F(RNTupleProcessorTest, Base)
 {
-   RNTupleOpenSpec ntuple{fNTupleName, fFileName};
+   RNTupleOpenSpec ntuple{fNTupleNames[0], fFileNames[0]};
    auto proc = RNTupleProcessor::Create(ntuple);
 
    int nEntries = 0;
@@ -64,7 +80,7 @@ TEST_F(RNTupleProcessorTest, Base)
 
 TEST_F(RNTupleProcessorTest, BaseWithModel)
 {
-   RNTupleOpenSpec ntuple{fNTupleName, fFileName};
+   RNTupleOpenSpec ntuple{fNTupleNames[0], fFileNames[0]};
 
    auto model = RNTupleModel::Create();
    auto fldX = model->MakeField<float>("x");
@@ -92,7 +108,7 @@ TEST_F(RNTupleProcessorTest, BaseWithModel)
 
 TEST_F(RNTupleProcessorTest, BaseWithBareModel)
 {
-   RNTupleOpenSpec ntuple{fNTupleName, fFileName};
+   RNTupleOpenSpec ntuple{fNTupleNames[0], fFileNames[0]};
 
    auto model = RNTupleModel::CreateBare();
    model->MakeField<float>("x");
@@ -117,3 +133,51 @@ TEST_F(RNTupleProcessorTest, BaseWithBareModel)
    EXPECT_EQ(nEntries, 5);
    EXPECT_EQ(nEntries, proc->GetNEntriesProcessed());
 }
+
+TEST_F(RNTupleProcessorTest, ChainedChain)
+{
+   std::vector<RNTupleOpenSpec> ntuples{{fNTupleNames[0], fFileNames[0]}, {fNTupleNames[0], fFileNames[0]}};
+
+   std::vector<std::unique_ptr<RNTupleProcessor>> innerProcs;
+   innerProcs.push_back(RNTupleProcessor::CreateChain(ntuples));
+   innerProcs.push_back(RNTupleProcessor::Create(ntuples[0]));
+
+   auto proc = RNTupleProcessor::CreateChain(std::move(innerProcs));
+
+   int nEntries = 0;
+
+   for (const auto &entry [[maybe_unused]] : *proc) {
+      EXPECT_EQ(++nEntries, proc->GetNEntriesProcessed());
+      EXPECT_EQ(nEntries - 1, proc->GetCurrentEntryNumber());
+      EXPECT_EQ(*entry.GetPtr<int>("i"), proc->GetCurrentEntryNumber() % 5);
+      EXPECT_EQ(static_cast<float>(*entry.GetPtr<int>("i")), *entry.GetPtr<float>("x"));
+   }
+   EXPECT_EQ(nEntries, 15);
+   EXPECT_EQ(nEntries, proc->GetNEntriesProcessed());
+}
+
+TEST_F(RNTupleProcessorTest, ChainedJoin)
+{
+   std::vector<RNTupleOpenSpec> ntuples{{fNTupleNames[0], fFileNames[0]}, {fNTupleNames[1], fFileNames[1]}};
+
+   std::vector<std::unique_ptr<RNTupleProcessor>> innerProcs;
+   innerProcs.push_back(RNTupleProcessor::CreateJoin(ntuples, {}));
+   innerProcs.push_back(RNTupleProcessor::CreateJoin(ntuples, {"i"}));
+
+   auto proc = RNTupleProcessor::CreateChain(std::move(innerProcs));
+
+   int nEntries = 0;
+
+   auto x = proc->GetEntry().GetPtr<float>("x");
+
+   for (const auto &entry [[maybe_unused]] : *proc) {
+      EXPECT_EQ(++nEntries, proc->GetNEntriesProcessed());
+      EXPECT_EQ(nEntries - 1, proc->GetCurrentEntryNumber());
+      EXPECT_EQ(*entry.GetPtr<int>("i"), proc->GetCurrentEntryNumber() % 5);
+
+      EXPECT_EQ(static_cast<float>(*entry.GetPtr<int>("i")), *x);
+      EXPECT_EQ(*x * 2, *entry.GetPtr<float>("ntuple_aux.z"));
+   }
+   EXPECT_EQ(nEntries, 10);
+   EXPECT_EQ(nEntries, proc->GetNEntriesProcessed());
+}