Skip to content

Commit 0aa713f

Browse files
enirolfenirolf
committed
[df] Expand FromTTree tests
Also check the resulting on-disk layout, as well as the effects of turning the `kVectorToRVec` snapshot option off.
1 parent bc3a50b commit 0aa713f

File tree

1 file changed

+145
-99
lines changed

1 file changed

+145
-99
lines changed

‎tree/dataframe/test/dataframe_snapshot_ntuple.cxx‎

Lines changed: 145 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -456,113 +456,159 @@ TEST(RDFSnapshotRNTuple, TDirectory)
456456
EXPECT_EQ(expected, sdf.GetColumnNames());
457457
}
458458

459-
void WriteTestTree(const std::string &tname, const std::string &fname)
460-
{
461-
TFile file(fname.c_str(), "RECREATE");
462-
TTree t(tname.c_str(), tname.c_str());
463-
464-
float pt = 42.f;
465-
std::vector<float> photons{1.f, 2.f, 3.f};
466-
Electron electron{137.f};
467-
Jet jets;
468-
jets.electrons.emplace_back(Electron{122.f});
469-
jets.electrons.emplace_back(Electron{125.f});
470-
jets.electrons.emplace_back(Electron{129.f});
471-
472-
Int_t nmuons = 1;
473-
float muon_pt[3] = {10.f, 20.f, 30.f};
474-
475-
struct {
476-
Int_t x = 1;
477-
Int_t y = 2;
478-
} point;
479-
480-
t.Branch("pt", &pt);
481-
t.Branch("photons", &photons);
482-
t.Branch("electron", &electron);
483-
t.Branch("jets", &jets);
484-
t.Branch("nmuons", &nmuons);
485-
t.Branch("muon_pt", muon_pt, "muon_pt[nmuons]");
486-
t.Branch("point", &point, "x/I:y/I");
487-
488-
t.Fill();
489-
t.Write();
490-
}
459+
class RDFSnapshotRNTupleFromTTreeTest : public ::testing::Test {
460+
protected:
461+
const std::string fFileName = "RDFSnapshotRNTuple_from_ttree.root";
462+
const std::string fTreeName = "tree";
491463

492-
TEST(RDFSnapshotRNTuple, FromTTree)
493-
{
494-
const auto treename = "tree";
495-
FileRAII fileGuard{"RDFSnapshotRNTuple_from_ttree.root"};
464+
void SetUp() override
465+
{
466+
TFile file(fFileName.c_str(), "RECREATE");
467+
TTree t(fTreeName.c_str(), fTreeName.c_str());
468+
469+
float pt = 42.f;
470+
std::vector<float> photons{1.f, 2.f, 3.f};
471+
Electron electron{137.f};
472+
Jet jets;
473+
jets.electrons.emplace_back(Electron{122.f});
474+
jets.electrons.emplace_back(Electron{125.f});
475+
jets.electrons.emplace_back(Electron{129.f});
476+
477+
Int_t nmuons = 1;
478+
float muon_pt[3] = {10.f, 20.f, 30.f};
479+
480+
struct {
481+
Int_t x = 1;
482+
Int_t y = 2;
483+
} point;
484+
485+
t.Branch("pt", &pt);
486+
t.Branch("photons", &photons);
487+
t.Branch("electron", &electron);
488+
t.Branch("jets", &jets);
489+
t.Branch("nmuons", &nmuons);
490+
t.Branch("muon_pt", muon_pt, "muon_pt[nmuons]");
491+
t.Branch("point", &point, "x/I:y/I");
492+
493+
t.Fill();
494+
t.Write();
495+
}
496496

497-
WriteTestTree(treename, fileGuard.GetPath());
497+
void TearDown() override
498+
{
499+
gSystem->Unlink(fFileName.c_str());
500+
std::remove(fFileName.c_str());
501+
}
498502

499-
auto df = ROOT::RDataFrame(treename, fileGuard.GetPath());
503+
void TestFromTTree(const std::string &fname, bool vector2RVec = false)
504+
{
505+
FileRAII fileGuard{fname};
500506

501-
RSnapshotOptions opts;
502-
opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
507+
auto df = ROOT::RDataFrame(fTreeName, fFileName);
503508

504-
{
505-
// FIXME(fdegeus): snapshotting leaflist branches as-is (i.e. without explicitly providing their leafs) is not
506-
// supported, because we have no way of reconstructing the memory layout of the branch itself from only the
507-
// TTree's on-disk information without JITting. For RNTuple, we would be able to do this using anonymous record
508-
// fields, however. Once this is implemented, this test should be changed to check the result of snapshotting
509-
// "point" fully.
510-
auto sdf = df.Define("x", [] { return 10; })
511-
.Snapshot("ntuple", fileGuard.GetPath(),
512-
{"x", "pt", "photons", "electron", "jets", "muon_pt", "point.x", "point.y"}, opts);
513-
514-
auto x = sdf->Take<int>("x");
515-
auto pt = sdf->Take<float>("pt");
516-
auto photons = sdf->Take<ROOT::RVec<float>>("photons");
517-
auto electron = sdf->Take<Electron>("electron");
518-
auto jet_electrons = sdf->Take<ROOT::RVec<Electron>>("jets.electrons");
519-
auto nMuons = sdf->Take<int>("nmuons");
520-
auto muonPt = sdf->Take<ROOT::RVec<float>>("muon_pt");
521-
auto pointX = sdf->Take<int>("point_x");
522-
auto pointY = sdf->Take<int>("point_y");
523-
524-
ASSERT_EQ(1UL, x->size());
525-
ASSERT_EQ(1UL, pt->size());
526-
ASSERT_EQ(1UL, photons->size());
527-
ASSERT_EQ(1UL, electron->size());
528-
ASSERT_EQ(1UL, jet_electrons->size());
529-
ASSERT_EQ(1UL, nMuons->size());
530-
ASSERT_EQ(1UL, muonPt->size());
531-
ASSERT_EQ(1UL, pointX->size());
532-
ASSERT_EQ(1UL, pointY->size());
533-
534-
EXPECT_EQ(10, x->front());
535-
EXPECT_EQ(42.f, pt->front());
536-
expect_vec_eq<float>({1.f, 2.f, 3.f}, photons->front());
537-
EXPECT_EQ(Electron{137.f}, electron->front());
538-
expect_vec_eq({Electron{122.f}, Electron{125.f}, Electron{129.f}}, jet_electrons->front());
539-
EXPECT_EQ(1, nMuons->front());
540-
expect_vec_eq({10.f}, muonPt->front());
541-
EXPECT_EQ(1, pointX->front());
542-
EXPECT_EQ(2, pointY->front());
543-
}
509+
{
510+
RSnapshotOptions opts;
511+
opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
512+
opts.fVector2RVec = vector2RVec;
513+
514+
// FIXME(fdegeus): snapshotting leaflist branches as-is (i.e. without explicitly providing their leafs) is not
515+
// supported, because we have no way of reconstructing the memory layout of the branch itself from only the
516+
// TTree's on-disk information without JITting. For RNTuple, we would be able to do this using anonymous record
517+
// fields, however. Once this is implemented, this test should be changed to check the result of snapshotting
518+
// "point" fully.
519+
auto sdf = df.Define("x", [] { return 10; })
520+
.Snapshot("ntuple", fileGuard.GetPath(),
521+
{"x", "pt", "photons", "electron", "jets", "muon_pt", "point.x", "point.y"}, opts);
522+
523+
auto x = sdf->Take<int>("x");
524+
auto pt = sdf->Take<float>("pt");
525+
auto photons = sdf->Take<ROOT::RVec<float>>("photons");
526+
auto electron = sdf->Take<Electron>("electron");
527+
auto jet_electrons = sdf->Take<ROOT::RVec<Electron>>("jets.electrons");
528+
auto nMuons = sdf->Take<int>("nmuons");
529+
auto muonPt = sdf->Take<ROOT::RVec<float>>("muon_pt");
530+
auto pointX = sdf->Take<int>("point_x");
531+
auto pointY = sdf->Take<int>("point_y");
532+
533+
ASSERT_EQ(1UL, x->size());
534+
ASSERT_EQ(1UL, pt->size());
535+
ASSERT_EQ(1UL, photons->size());
536+
ASSERT_EQ(1UL, electron->size());
537+
ASSERT_EQ(1UL, jet_electrons->size());
538+
ASSERT_EQ(1UL, nMuons->size());
539+
ASSERT_EQ(1UL, muonPt->size());
540+
ASSERT_EQ(1UL, pointX->size());
541+
ASSERT_EQ(1UL, pointY->size());
542+
543+
EXPECT_EQ(10, x->front());
544+
EXPECT_EQ(42.f, pt->front());
545+
expect_vec_eq<float>({1.f, 2.f, 3.f}, photons->front());
546+
EXPECT_EQ(Electron{137.f}, electron->front());
547+
expect_vec_eq({Electron{122.f}, Electron{125.f}, Electron{129.f}}, jet_electrons->front());
548+
EXPECT_EQ(1, nMuons->front());
549+
expect_vec_eq({10.f}, muonPt->front());
550+
EXPECT_EQ(1, pointX->front());
551+
EXPECT_EQ(2, pointY->front());
552+
}
544553

545-
auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
554+
auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
546555

547-
auto x = reader->GetView<int>("x");
548-
auto pt = reader->GetView<float>("pt");
549-
auto photons = reader->GetView<ROOT::RVec<float>>("photons");
550-
auto electron = reader->GetView<Electron>("electron");
551-
auto jet_electrons = reader->GetView<ROOT::RVec<Electron>>("jets.electrons");
552-
auto nMuons = reader->GetView<int>("nmuons");
553-
auto muonPt = reader->GetView<ROOT::RVec<float>>("muon_pt");
554-
auto pointX = reader->GetView<int>("point_x");
555-
auto pointY = reader->GetView<int>("point_y");
556+
auto &descriptor = reader->GetDescriptor();
556557

557-
EXPECT_EQ(10, x(0));
558-
EXPECT_EQ(42.f, pt(0));
559-
expect_vec_eq<float>({1.f, 2.f, 3.f}, photons(0));
560-
EXPECT_EQ(Electron{137.f}, electron(0));
561-
expect_vec_eq({Electron{122.f}, Electron{125.f}, Electron{129.f}}, jet_electrons(0));
562-
EXPECT_EQ(1, nMuons(0));
563-
expect_vec_eq({10.f}, muonPt(0));
564-
EXPECT_EQ(1, pointX(0));
565-
EXPECT_EQ(2, pointY(0));
558+
int nTopLevelFields = std::distance(descriptor.GetTopLevelFields().begin(), descriptor.GetTopLevelFields().end());
559+
EXPECT_EQ(9, nTopLevelFields);
560+
EXPECT_EQ("std::int32_t", descriptor.GetFieldDescriptor(descriptor.FindFieldId("x")).GetTypeName());
561+
EXPECT_EQ("float", descriptor.GetFieldDescriptor(descriptor.FindFieldId("pt")).GetTypeName());
562+
if (vector2RVec) {
563+
EXPECT_EQ("ROOT::VecOps::RVec<float>",
564+
descriptor.GetFieldDescriptor(descriptor.FindFieldId("photons")).GetTypeName());
565+
} else {
566+
EXPECT_EQ("std::vector<float>",
567+
descriptor.GetFieldDescriptor(descriptor.FindFieldId("photons")).GetTypeName());
568+
}
569+
EXPECT_EQ("Electron", descriptor.GetFieldDescriptor(descriptor.FindFieldId("electron")).GetTypeName());
570+
auto jetsId = descriptor.FindFieldId("jets");
571+
EXPECT_EQ("Jet", descriptor.GetFieldDescriptor(jetsId).GetTypeName());
572+
auto electronsId = descriptor.FindFieldId("electrons", jetsId);
573+
EXPECT_EQ("std::vector<Electron>", descriptor.GetFieldDescriptor(electronsId).GetTypeName());
574+
EXPECT_EQ("std::int32_t", descriptor.GetFieldDescriptor(descriptor.FindFieldId("nmuons")).GetTypeName());
575+
EXPECT_EQ("ROOT::VecOps::RVec<float>",
576+
descriptor.GetFieldDescriptor(descriptor.FindFieldId("muon_pt")).GetTypeName());
577+
EXPECT_EQ("std::int32_t", descriptor.GetFieldDescriptor(descriptor.FindFieldId("point_x")).GetTypeName());
578+
EXPECT_EQ("std::int32_t", descriptor.GetFieldDescriptor(descriptor.FindFieldId("point_y")).GetTypeName());
579+
// sanity check to make sure we don't snapshot the internal RDF size columns
580+
EXPECT_EQ(ROOT::kInvalidDescriptorId, descriptor.FindFieldId("R_rdf_sizeof_photons"));
581+
582+
auto x = reader->GetView<int>("x");
583+
auto pt = reader->GetView<float>("pt");
584+
auto photons = reader->GetView<ROOT::RVec<float>>("photons");
585+
auto electron = reader->GetView<Electron>("electron");
586+
auto jet_electrons = reader->GetView<ROOT::RVec<Electron>>("jets.electrons");
587+
auto nMuons = reader->GetView<int>("nmuons");
588+
auto muonPt = reader->GetView<ROOT::RVec<float>>("muon_pt");
589+
auto pointX = reader->GetView<int>("point_x");
590+
auto pointY = reader->GetView<int>("point_y");
591+
592+
EXPECT_EQ(10, x(0));
593+
EXPECT_EQ(42.f, pt(0));
594+
expect_vec_eq<float>({1.f, 2.f, 3.f}, photons(0));
595+
EXPECT_EQ(Electron{137.f}, electron(0));
596+
expect_vec_eq({Electron{122.f}, Electron{125.f}, Electron{129.f}}, jet_electrons(0));
597+
EXPECT_EQ(1, nMuons(0));
598+
expect_vec_eq({10.f}, muonPt(0));
599+
EXPECT_EQ(1, pointX(0));
600+
EXPECT_EQ(2, pointY(0));
601+
}
602+
};
603+
604+
TEST_F(RDFSnapshotRNTupleFromTTreeTest, FromTTree)
605+
{
606+
TestFromTTree("RDFSnapshotRNTuple_from_ttree.root");
607+
}
608+
609+
TEST_F(RDFSnapshotRNTupleFromTTreeTest, FromTTreeNoVector2RVec)
610+
{
611+
TestFromTTree("RDFSnapshotRNTuple_from_ttree_novec2rvec.root", /*vector2RVec=*/false);
566612
}
567613

568614
#ifdef R__USE_IMT

0 commit comments

Comments
 (0)