Skip to content

Commit 3e1ab31

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 3e1ab31

File tree

1 file changed

+141
-99
lines changed

1 file changed

+141
-99
lines changed

‎tree/dataframe/test/dataframe_snapshot_ntuple.cxx‎

Lines changed: 141 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -456,113 +456,155 @@ 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_ttree_fixture.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 { gSystem->Unlink(fFileName.c_str()); }
498498

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

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

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

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

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");
552+
auto &descriptor = reader->GetDescriptor();
556553

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

568610
#ifdef R__USE_IMT

0 commit comments

Comments
 (0)