[ntuple] Add tests for TMemFile-based reading

Author: enirolf

tree/ntuple/v7/test/ntuple_processor.cxx

@@ -2,6 +2,8 @@
 
 #include <ROOT/RNTupleProcessor.hxx>
 
+#include <TMemFile.h>
+
 TEST(RNTupleProcessor, EmptyNTuple)
 {
    FileRaii fileGuard("test_ntuple_processor_empty.root");
@@ -22,6 +24,35 @@ TEST(RNTupleProcessor, EmptyNTuple)
    EXPECT_EQ(nEntries, proc->GetNEntriesProcessed());
 }
 
+TEST(RNTupleProcessor, TMemFile)
+{
+   {
+      auto memFile = std::make_unique<TMemFile>("ntuple", "RECREATE");
+      auto model = RNTupleModel::Create();
+      auto fldX = model->MakeField<float>("x");
+      auto ntuple = RNTupleWriter::Append(std::move(model), "ntuple", *memFile);
+
+      for (unsigned i = 0; i < 5; ++i) {
+         *fldX = static_cast<float>(i);
+         ntuple->Fill();
+      }
+   }
+
+   auto memFile = std::make_unique<TMemFile>("ntuple");
+   auto proc = RNTupleProcessor::Create({"ntuple", memFile.get()});
+   auto x = proc->GetEntry().GetPtr<float>("x");
+
+   int nEntries = 0;
+   for ([[maybe_unused]] const auto &entry : *proc) {
+      EXPECT_EQ(++nEntries, proc->GetNEntriesProcessed());
+      EXPECT_EQ(nEntries - 1, proc->GetCurrentEntryNumber());
+
+      EXPECT_FLOAT_EQ(static_cast<float>(nEntries - 1), *x);
+   }
+   EXPECT_EQ(nEntries, 5);
+   EXPECT_EQ(nEntries, proc->GetNEntriesProcessed());
+}
+
 class RNTupleProcessorTest : public testing::Test {
 protected:
    const std::array<std::string, 3> fFileNames{"test_ntuple_processor1.root", "test_ntuple_processor2.root",

tree/ntuple/v7/test/ntuple_processor_chain.cxx

@@ -2,6 +2,8 @@
 
 #include <ROOT/RNTupleProcessor.hxx>
 
+#include <TMemFile.h>
+
 class RNTupleChainProcessorTest : public testing::Test {
 protected:
    const std::array<std::string, 4> fFileNames{"test_ntuple_chain_processor1.root", "test_ntuple_chain_processor2.root",
@@ -258,3 +260,44 @@ TEST_F(RNTupleChainProcessorTest, LoadRandomEntry)
 
    EXPECT_EQ(ROOT::kInvalidNTupleIndex, RNTupleProcessorEntryLoader::LoadEntry(*proc, 8));
 }
+
+TEST(RNTupleChainProcessor, TMemFile)
+{
+   FileRaii fileGuard("test_ntuple_processor_join_tmemfile.root");
+   {
+      auto model = RNTupleModel::Create();
+      auto fldX = model->MakeField<float>("x");
+      auto ntuple = RNTupleWriter::Recreate(std::move(model), "ntuple", fileGuard.GetPath());
+
+      for (unsigned i = 0; i < 5; i++) {
+         *fldX = static_cast<float>(i);
+         ntuple->Fill();
+      }
+   }
+   {
+      auto memFile = std::make_unique<TMemFile>("ntuple", "RECREATE");
+      auto model = RNTupleModel::Create();
+      auto fldX = model->MakeField<float>("x");
+      auto ntuple = RNTupleWriter::Append(std::move(model), "ntuple", *memFile);
+
+      for (unsigned i = 5; i < 10; ++i) {
+         *fldX = static_cast<float>(i);
+         ntuple->Fill();
+      }
+   }
+
+   auto memFile = std::make_unique<TMemFile>("ntuple");
+
+   auto proc = RNTupleProcessor::CreateChain({{"ntuple", fileGuard.GetPath()}, {"ntuple", memFile.get()}});
+
+   std::uint64_t nEntries = 0;
+   auto x = proc->GetEntry().GetPtr<float>("x");
+   for ([[maybe_unused]] const auto &entry : *proc) {
+      EXPECT_EQ(++nEntries, proc->GetNEntriesProcessed());
+      EXPECT_EQ(nEntries - 1, proc->GetCurrentEntryNumber());
+
+      EXPECT_EQ(static_cast<float>(nEntries - 1), *x);
+   }
+   EXPECT_EQ(nEntries, 10);
+   EXPECT_EQ(nEntries, proc->GetNEntriesProcessed());
+}

tree/ntuple/v7/test/ntuple_processor_join.cxx

@@ -2,6 +2,8 @@
 
 #include <ROOT/RNTupleProcessor.hxx>
 
+#include <TMemFile.h>
+
 class RNTupleJoinProcessorTest : public testing::Test {
 protected:
    const std::array<std::string, 4> fFileNames{"test_ntuple_join_processor1.root", "test_ntuple_join_processor2.root",
@@ -287,3 +289,53 @@ TEST_F(RNTupleJoinProcessorTest, WithModel)
 
    EXPECT_EQ(5, proc->GetNEntriesProcessed());
 }
+
+TEST(RNTupleJoinProcessor, TMemFile)
+{
+   FileRaii fileGuard("test_ntuple_processor_join_tmemfile.root");
+   {
+      auto model = RNTupleModel::Create();
+      auto fldI = model->MakeField<int>("i");
+      auto fldX = model->MakeField<float>("x");
+      auto ntuple = RNTupleWriter::Recreate(std::move(model), "ntuple", fileGuard.GetPath());
+
+      for (unsigned i = 0; i < 5; ++i) {
+         *fldI = i * 2;
+         *fldX = *fldI * 0.5f;
+         ntuple->Fill();
+      }
+   }
+   {
+      auto memFile = std::make_unique<TMemFile>("ntuple_aux", "RECREATE");
+      auto model = RNTupleModel::Create();
+      auto fldI = model->MakeField<int>("i");
+      auto fldY = model->MakeField<std::vector<float>>("y");
+      auto ntuple = RNTupleWriter::Append(std::move(model), "ntuple_aux", *memFile);
+
+      for (unsigned i = 0; i < 10; ++i) {
+         *fldI = i;
+         *fldY = {static_cast<float>(*fldI * 0.2), 3.14, static_cast<float>(*fldI * 1.3)};
+         ntuple->Fill();
+      }
+   }
+
+   auto memFile = std::make_unique<TMemFile>("ntuple_aux");
+   auto proc = RNTupleProcessor::CreateJoin({"ntuple", fileGuard.GetPath()}, {{"ntuple_aux", memFile.get()}}, {"i"});
+
+   int nEntries = 0;
+   auto i = proc->GetEntry().GetPtr<int>("i");
+   auto x = proc->GetEntry().GetPtr<float>("x");
+   auto y = proc->GetEntry().GetPtr<std::vector<float>>("ntuple_aux.y");
+   std::vector<float> yExpected;
+   for ([[maybe_unused]] auto &entry : *proc) {
+      EXPECT_EQ(proc->GetCurrentEntryNumber(), nEntries++);
+
+      EXPECT_FLOAT_EQ(proc->GetCurrentEntryNumber() * 2, *i);
+      EXPECT_FLOAT_EQ(*i * 0.5f, *x);
+
+      yExpected = {static_cast<float>(*i * 0.2), 3.14, static_cast<float>(*i * 1.3)};
+      EXPECT_EQ(yExpected, *y);
+   }
+
+   EXPECT_EQ(5, proc->GetNEntriesProcessed());
+}