Skip to content

Commit

Permalink
[ntuple] Add support for std::unordered_set fields
Browse files Browse the repository at this point in the history
  • Loading branch information
enirolf committed Nov 16, 2023
1 parent 19b6e6e commit ab04424
Show file tree
Hide file tree
Showing 5 changed files with 109 additions and 1 deletion.
27 changes: 26 additions & 1 deletion tree/ntuple/v7/inc/ROOT/RField.hxx
Original file line number Diff line number Diff line change
Expand Up @@ -1141,7 +1141,7 @@ public:
size_t GetAlignment() const final { return fMaxAlignment; }
};

/// The generic field for a std::set<Type>
/// The generic field for a std::set<Type> and std::unordered_set<Type>
class RSetField : public RProxiedCollectionField {
protected:
std::unique_ptr<Detail::RFieldBase> CloneImpl(std::string_view newName) const final;
Expand Down Expand Up @@ -2238,6 +2238,31 @@ public:
size_t GetAlignment() const final { return std::alignment_of<ContainerT>(); }
};

template <typename ItemT>
class RField<std::unordered_set<ItemT>> : public RSetField {
using ContainerT = typename std::unordered_set<ItemT>;

protected:
void GenerateValue(void *where) const final { new (where) ContainerT(); }
void DestroyValue(void *objPtr, bool dtorOnly = false) const final
{
std::destroy_at(static_cast<ContainerT *>(objPtr));
Detail::RFieldBase::DestroyValue(objPtr, dtorOnly);
}

public:
static std::string TypeName() { return "std::unordered_set<" + RField<ItemT>::TypeName() + ">"; }

explicit RField(std::string_view name) : RSetField(name, TypeName(), std::make_unique<RField<ItemT>>("_0")) {}
RField(RField &&other) = default;
RField &operator=(RField &&other) = default;
~RField() override = default;

using Detail::RFieldBase::GenerateValue;
size_t GetValueSize() const final { return sizeof(ContainerT); }
size_t GetAlignment() const final { return std::alignment_of<ContainerT>(); }
};

template <typename... ItemTs>
class RField<std::variant<ItemTs...>> : public RVariantField {
using ContainerT = typename std::variant<ItemTs...>;
Expand Down
8 changes: 8 additions & 0 deletions tree/ntuple/v7/src/RField.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,8 @@ std::string GetNormalizedTypeName(const std::string &typeName)
normalizedType = "std::" + normalizedType;
if (normalizedType.substr(0, 4) == "set<")
normalizedType = "std::" + normalizedType;
if (normalizedType.substr(0, 14) == "unordered_set<")
normalizedType = "std::" + normalizedType;
if (normalizedType.substr(0, 7) == "atomic<")
normalizedType = "std::" + normalizedType;
if (normalizedType == "byte")
Expand Down Expand Up @@ -456,6 +458,12 @@ ROOT::Experimental::Detail::RFieldBase::Create(const std::string &fieldName, con
auto normalizedInnerTypeName = itemField->GetType();
result =
std::make_unique<RSetField>(fieldName, "std::set<" + normalizedInnerTypeName + ">", std::move(itemField));
} else if (canonicalType.substr(0, 19) == "std::unordered_set<") {
std::string itemTypeName = canonicalType.substr(19, canonicalType.length() - 20);
auto itemField = Create("_0", itemTypeName).Unwrap();
auto normalizedInnerTypeName = itemField->GetType();
result = std::make_unique<RSetField>(fieldName, "std::unordered_set<" + normalizedInnerTypeName + ">",
std::move(itemField));
} else if (canonicalType.substr(0, 12) == "std::atomic<") {
std::string itemTypeName = canonicalType.substr(12, canonicalType.length() - 13);
auto itemField = Create("_0", itemTypeName).Unwrap();
Expand Down
6 changes: 6 additions & 0 deletions tree/ntuple/v7/test/CustomStruct.hxx
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include <cstdint>
#include <set>
#include <string>
#include <unordered_set>
#include <variant>
#include <vector>

Expand Down Expand Up @@ -38,6 +39,11 @@ struct CustomStruct {
bool operator==(const CustomStruct &c) const { return a == c.a && v1 == c.v1 && v2 == c.v2 && s == c.s; }
};

template <>
struct std::hash<CustomStruct> {
std::size_t operator()(const CustomStruct &c) const noexcept { return std::hash<float>{}(c.a); }
};

struct DerivedA : public CustomStruct {
std::vector<float> a_v;
std::string a_s;
Expand Down
6 changes: 6 additions & 0 deletions tree/ntuple/v7/test/CustomStructLinkDef.h
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,12 @@
#pragma link C++ class std::set<std::set<char>> +;
#pragma link C++ class std::set<std::pair<int, CustomStruct>> +;

#pragma link C++ class std::unordered_set<std::int64_t> +;
#pragma link C++ class std::unordered_set<std::string> +;
#pragma link C++ class std::unordered_set<float> +;
#pragma link C++ class std::unordered_set<CustomStruct> +;
#pragma link C++ class std::unordered_set<std::vector<bool>> +;

#pragma link C++ options = version(3) class StructWithIORulesBase + ;
#pragma link C++ options = version(3) class StructWithTransientString + ;
#pragma link C++ options = version(3) class StructWithIORules + ;
Expand Down
63 changes: 63 additions & 0 deletions tree/ntuple/v7/test/ntuple_types.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,69 @@ TEST(RNTuple, StdSet)
EXPECT_EQ(pairSet, *mySet2);
}

TEST(RNTuple, StdUnorderedSet)
{
auto field = RField<std::unordered_set<int64_t>>("setField");
EXPECT_STREQ("std::unordered_set<std::int64_t>", field.GetType().c_str());
auto otherField = RFieldBase::Create("test", "std::unordered_set<int64_t>").Unwrap();
EXPECT_STREQ(field.GetType().c_str(), otherField->GetType().c_str());
EXPECT_EQ((sizeof(std::unordered_set<int64_t>)), field.GetValueSize());
EXPECT_EQ((sizeof(std::unordered_set<int64_t>)), otherField->GetValueSize());
EXPECT_EQ((alignof(std::unordered_set<int64_t>)), field.GetAlignment());
// For type-erased set fields, we use `alignof(std::set<std::max_align_t>)` to set the alignment,
// so the actual alignment may be smaller.
EXPECT_LE((alignof(std::unordered_set<int64_t>)), otherField->GetAlignment());

FileRaii fileGuard("test_ntuple_rfield_stdunorderedset.root");
{
auto model = RNTupleModel::Create();
auto set_field = model->MakeField<std::unordered_set<float>>({"mySet", "unordered float set"});
auto set_field2 = model->MakeField<std::unordered_set<CustomStruct>>({"mySet2"});

auto mySet3 = RFieldBase::Create("mySet3", "std::unordered_set<std::string>").Unwrap();
auto mySet4 = RFieldBase::Create("mySet4", "std::unordered_set<std::vector<bool>>").Unwrap();

model->AddField(std::move(mySet3));
model->AddField(std::move(mySet4));

auto ntuple = RNTupleWriter::Recreate(std::move(model), "set_ntuple", fileGuard.GetPath());
auto set_field3 = ntuple->GetModel()->GetDefaultEntry()->Get<std::unordered_set<std::string>>("mySet3");
auto set_field4 = ntuple->GetModel()->GetDefaultEntry()->Get<std::unordered_set<std::vector<bool>>>("mySet4");
for (int i = 0; i < 2; i++) {
*set_field = {static_cast<float>(i), 3.14, 0.42};
*set_field2 = {CustomStruct{6.f, {7.f, 8.f}, {{9.f}, {10.f}}, "foo"},
CustomStruct{2.f, {3.f, 4.f}, {{5.f}, {6.f}}, "bar"}};
*set_field3 = {"Hello", "world!", std::to_string(i)};
*set_field4 = {{(i % 2 == 0)}, {}, {false, true}};
ntuple->Fill();
}
}

auto ntuple = RNTupleReader::Open("set_ntuple", fileGuard.GetPath());
EXPECT_EQ(2, ntuple->GetNEntries());

auto viewSet = ntuple->GetView<std::unordered_set<float>>("mySet");
auto viewSet2 = ntuple->GetView<std::unordered_set<CustomStruct>>("mySet2");
auto viewSet3 = ntuple->GetView<std::unordered_set<std::string>>("mySet3");
auto viewSet4 = ntuple->GetView<std::unordered_set<std::vector<bool>>>("mySet4");
for (auto i : ntuple->GetEntryRange()) {
EXPECT_EQ(std::unordered_set<float>({static_cast<float>(i), 3.14, 0.42}), viewSet(i));

auto pairSet = std::unordered_set<CustomStruct>(
{CustomStruct{6.f, {7.f, 8.f}, {{9.f}, {10.f}}, "foo"}, CustomStruct{2.f, {3.f, 4.f}, {{5.f}, {6.f}}, "bar"}});
EXPECT_EQ(pairSet, viewSet2(i));

EXPECT_EQ(std::unordered_set<std::string>({"Hello", "world!", std::to_string(i)}), viewSet3(i));
EXPECT_EQ(std::unordered_set<std::vector<bool>>({{(i % 2 == 0)}, {}, {false, true}}), viewSet4(i));
}

ntuple->LoadEntry(0);
auto mySet2 = ntuple->GetModel()->GetDefaultEntry()->Get<std::unordered_set<CustomStruct>>("mySet2");
auto pairSet = std::unordered_set<CustomStruct>(
{CustomStruct{6.f, {7.f, 8.f}, {{9.f}, {10.f}}, "foo"}, CustomStruct{2.f, {3.f, 4.f}, {{5.f}, {6.f}}, "bar"}});
EXPECT_EQ(pairSet, *mySet2);
}

TEST(RNTuple, Int64)
{
auto field = RFieldBase::Create("test", "std::int64_t").Unwrap();
Expand Down

0 comments on commit ab04424

Please sign in to comment.