Skip to content

[ntuple] Fix zero pages for encoded and deferred columns #23742

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Elvand-Lie wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[ntuple] Fix zero pages for encoded and deferred columns #23742

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 60 additions & 49 deletions tree/ntuple/src/RNTupleMerger.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1137,6 +1137,21 @@ RColumnInfoGroup RNTupleMergeStrategyDefault::GatherColumnInfos(const RDescripto
return res;
}

// Collection and variant children have no default elements for a missing entry.
static std::uint64_t GetElementsPerEntry(const ROOT::RFieldDescriptor &field, const ROOT::RNTupleDescriptor &descriptor)
{
auto nElements = std::max<std::uint64_t>(field.GetNRepetitions(), 1);
for (auto parentId = field.GetParentId(); parentId != ROOT::kInvalidDescriptorId;) {
const auto &parent = descriptor.GetFieldDescriptor(parentId);
if (parent.GetStructure() == ROOT::ENTupleStructure::kCollection ||
parent.GetStructure() == ROOT::ENTupleStructure::kVariant)
return 0;
nElements *= std::max<std::uint64_t>(parent.GetNRepetitions(), 1);
parentId = parent.GetParentId();
}
return nElements;
}

// Generates default (zero) values for the given columns
ROOT::RResult<void> RNTupleMergeStrategyDefault::GenerateZeroPagesForColumns(size_t nEntriesToGenerate,
std::span<const RColumnMergeInfo> columns,
Expand All @@ -1153,21 +1168,8 @@ ROOT::RResult<void> RNTupleMergeStrategyDefault::GenerateZeroPagesForColumns(siz
if (field->GetLogicalColumnIds()[0] != column.fInputId)
continue;

// Check if this column is a child of a Collection or a Variant. If so, it has no data
// and can be skipped.
bool skipColumn = false;
auto nRepetitions = std::max<std::uint64_t>(field->GetNRepetitions(), 1);
for (auto parentId = field->GetParentId(); parentId != ROOT::kInvalidDescriptorId;) {
const ROOT::RFieldDescriptor &parent = column.fParentNTupleDescriptor->GetFieldDescriptor(parentId);
if (parent.GetStructure() == ROOT::ENTupleStructure::kCollection ||
parent.GetStructure() == ROOT::ENTupleStructure::kVariant) {
skipColumn = true;
break;
}
nRepetitions *= std::max<std::uint64_t>(parent.GetNRepetitions(), 1);
parentId = parent.GetParentId();
}
if (skipColumn)
const auto nElementsPerEntry = GetElementsPerEntry(*field, *column.fParentNTupleDescriptor);
if (!nElementsPerEntry)
continue;

const auto structure = field->GetStructure();
Expand All @@ -1184,35 +1186,40 @@ ROOT::RResult<void> RNTupleMergeStrategyDefault::GenerateZeroPagesForColumns(siz
structure == ROOT::ENTupleStructure::kPlain);

const auto &columnDesc = fDstDescriptor.GetColumnDescriptor(column.fOutputId);
const auto colElement = RColumnElementBase::Generate(columnDesc.GetType());
const auto nElements = nEntriesToGenerate * nRepetitions;
const auto nBytesOnStorage = colElement->GetPackedSize(nElements);
// TODO(gparolini): make this configurable
constexpr auto kPageSizeLimit = 256 * 1024;
// TODO(gparolini): consider coalescing the last page if its size is less than some threshold
const size_t nPages = nBytesOnStorage / kPageSizeLimit + !!(nBytesOnStorage % kPageSizeLimit);
for (size_t i = 0; i < nPages; ++i) {
const auto pageSize = (i < nPages - 1) ? kPageSizeLimit : nBytesOnStorage - kPageSizeLimit * (nPages - 1);
const auto checksumSize = RPageStorage::kNBytesPageChecksum;
const auto bufSize = pageSize + checksumSize;
assert(pageSize % colElement->GetSize() == 0);
const auto nElementsPerPage = pageSize / colElement->GetSize();
auto page = fPageAlloc->NewPage(colElement->GetSize(), nElementsPerPage);
page.GrowUnchecked(nElementsPerPage);
memset(page.GetBuffer(), 0, page.GetNBytes());

auto &buffer = sealedPageData.fBuffers.emplace_back(new std::byte[bufSize]);
RPageSink::RSealPageConfig sealConf;
sealConf.fElement = colElement.get();
sealConf.fPage = &page;
sealConf.fBuffer = buffer.get();
sealConf.fCompressionSettings = fMergeOpts.fCompressionSettings.value();
sealConf.fWriteChecksum = fDestination.GetWriteOptions().GetEnablePageChecksums();
auto sealedPage = RPageSink::SealPage(sealConf);

sealedPageData.fPagesV.push_back({sealedPage});
sealedPageData.fGroups.emplace_back(column.fOutputId, sealedPageData.fPagesV.back().cbegin(),
sealedPageData.fPagesV.back().cend());
try {
const auto colElement = RColumnElementBase::Generate(columnDesc.GetType());
colElement->SetBitsOnStorage(columnDesc.GetBitsOnStorage());
if (const auto &range = columnDesc.GetValueRange(); range)
colElement->SetValueRange(range->fMin, range->fMax);
const auto nElements = nEntriesToGenerate * nElementsPerEntry;
// TODO(gparolini): make this configurable
constexpr auto kPageSizeLimit = 256 * 1024;
const auto maxElementsPerPage = std::max<std::size_t>(kPageSizeLimit * 8 / colElement->GetBitsOnStorage(), 1);
// TODO(gparolini): consider coalescing the last page if its size is less than some threshold
for (std::size_t nRemaining = nElements; nRemaining > 0;) {
const auto nElementsPerPage = std::min(nRemaining, maxElementsPerPage);
const auto bufSize = colElement->GetPackedSize(nElementsPerPage) + RPageStorage::kNBytesPageChecksum;
auto page = fPageAlloc->NewPage(colElement->GetSize(), nElementsPerPage);
page.GrowUnchecked(nElementsPerPage);
memset(page.GetBuffer(), 0, page.GetNBytes());

auto &buffer = sealedPageData.fBuffers.emplace_back(new std::byte[bufSize]);
RPageSink::RSealPageConfig sealConf;
sealConf.fElement = colElement.get();
sealConf.fPage = &page;
sealConf.fBuffer = buffer.get();
sealConf.fCompressionSettings = fMergeOpts.fCompressionSettings.value();
sealConf.fWriteChecksum = fDestination.GetWriteOptions().GetEnablePageChecksums();
auto sealedPage = RPageSink::SealPage(sealConf);

sealedPageData.fPagesV.push_back({sealedPage});
sealedPageData.fGroups.emplace_back(column.fOutputId, sealedPageData.fPagesV.back().cbegin(),
sealedPageData.fPagesV.back().cend());
nRemaining -= nElementsPerPage;
}
} catch (const ROOT::RException &ex) {
return R__FAIL("cannot generate default values for field `" +
column.fParentNTupleDescriptor->GetQualifiedFieldName(field->GetId()) + "`: " + ex.what());
}
}
return ROOT::RResult<void>::Success();
Expand Down Expand Up @@ -1287,11 +1294,15 @@ ROOT::RResult<void> RNTupleMergeStrategyDefault::MergeCommonColumns(ROOT::Intern
// As an optimization, we don't do this for the first source (since we can rely on the FirstElementIndex and
// deferred column mechanism in that case).
// TODO: also avoid doing this if we added no real page of this column to the destination yet.
if (columnDesc.GetFirstElementIndex() > clusterDesc.GetFirstEntryIndex() && fDestination.GetNEntries() > 0) {
const auto nMissingEntries = columnDesc.GetFirstElementIndex() - clusterDesc.GetFirstEntryIndex();
auto res = GenerateZeroPagesForColumns(nMissingEntries, {&column, 1}, sealedPageData);
if (!res)
return R__FORWARD_ERROR(res);
if (const auto nElementsPerEntry =
GetElementsPerEntry(*column.fParentFieldDescriptor, *column.fParentNTupleDescriptor)) {
const auto firstEntry = columnDesc.GetFirstElementIndex() / nElementsPerEntry;
if (firstEntry > clusterDesc.GetFirstEntryIndex() && fDestination.GetNEntries() > 0) {
const auto nMissingEntries = firstEntry - clusterDesc.GetFirstEntryIndex();
auto res = GenerateZeroPagesForColumns(nMissingEntries, {&column, 1}, sealedPageData);
if (!res)
return R__FORWARD_ERROR(res);
}
}

// Loop over the pages
Expand Down
219 changes: 219 additions & 0 deletions tree/ntuple/test/ntuple_merger.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@
#include <TRandom3.h>
#include <TTree.h>

#include <bitset>
#include <string_view>
#include <type_traits>
#include <unordered_map>

using ROOT::TestSupport::CheckDiagsRAII;
Expand Down Expand Up @@ -2454,6 +2456,223 @@ TEST(RNTupleMerger, SkipMissing)
EXPECT_TRUE(ok);
}

struct RNTupleMergerZeroPages : public ::testing::TestWithParam<std::string> {};

TEST_P(RNTupleMergerZeroPages, MissingField)
{
const auto encoding = GetParam();
auto check = [&](auto sample) {
using T = decltype(sample);
FileRaii first("test_ntuple_merge_zero_pages_first.root");
FileRaii middle("test_ntuple_merge_zero_pages_middle.root");
FileRaii last("test_ntuple_merge_zero_pages_last.root");
FileRaii output("test_ntuple_merge_zero_pages_output.root");
const std::array<const FileRaii *, 3> files{&first, &middle, &last};
constexpr unsigned int nEntriesPerSource = 20;
for (unsigned int source = 0; source < files.size(); ++source) {
auto model = RNTupleModel::Create();
auto entry = model->MakeField<unsigned int>("entry");
if (source != 1) {
auto field = std::make_unique<RField<T>>("value");
if constexpr (std::is_same_v<T, float>) {
if (encoding == "Real16")
field->SetHalfPrecision();
else if (encoding == "Real32Trunc")
field->SetTruncated(20);
else
field->SetQuantized(20, {0., 64.});
}
model->AddField(std::move(field));
}
auto writer = RNTupleWriter::Recreate(std::move(model), "ntuple", files[source]->GetPath());
std::shared_ptr<T> value;
if (source != 1)
value = writer->GetModel().GetDefaultEntry().template GetPtr<T>("value");
for (unsigned int i = 0; i < nEntriesPerSource; ++i) {
*entry = source * nEntriesPerSource + i;
if (value)
*value = T(*entry + 1);
writer->Fill();
}
}

std::vector<std::unique_ptr<RPageSource>> sources;
std::vector<RPageSource *> sourcePtrs;
for (const auto *file : files) {
sources.push_back(RPageSource::Create("ntuple", file->GetPath()));
sourcePtrs.push_back(sources.back().get());
}
auto destination = std::make_unique<RPageSinkFile>("ntuple", output.GetPath(), RNTupleWriteOptions());
RNTupleMerger merger{std::move(destination)};
RNTupleMergeOptions options;
options.fMergingMode = ENTupleMergingMode::kUnion;
auto result = merger.Merge(sourcePtrs, options);
ASSERT_TRUE(bool(result)) << result.GetError()->GetReport();

auto reader = RNTupleReader::Open("ntuple", output.GetPath());
ASSERT_EQ(3 * nEntriesPerSource, reader->GetNEntries());
const auto &descriptor = reader->GetDescriptor();
const auto &field = descriptor.GetFieldDescriptor(descriptor.FindFieldId("value"));
ASSERT_EQ(1u, field.GetLogicalColumnIds().size());
const auto columnId = field.GetLogicalColumnIds()[0];
const auto repetitions = std::max<std::uint64_t>(field.GetNRepetitions(), 1);
ASSERT_EQ(3u, descriptor.GetNActiveClusters());
for (const auto &cluster : descriptor.GetActiveClusterIterable()) {
ASSERT_EQ(cluster.GetNEntries() * repetitions, cluster.GetColumnRange(columnId).GetNElements());
}
auto value = reader->GetModel().GetDefaultEntry().template GetPtr<T>("value");
for (auto i : reader->GetEntryRange()) {
SCOPED_TRACE(i);
reader->LoadEntry(i);
const auto expected = (i >= nEntriesPerSource && i < 2 * nEntriesPerSource) ? T{} : T(i + 1);
if constexpr (std::is_same_v<T, float>)
EXPECT_NEAR(expected, *value, 0.001f);
else
EXPECT_EQ(expected, *value);
}
};
if (encoding == "Bool")
check(bool{});
else if (encoding == "Bitset8")
check(std::bitset<8>{});
else
check(float{});
}

INSTANTIATE_TEST_SUITE_P(Encoded, RNTupleMergerZeroPages,
::testing::Values("Bool", "Bitset8", "Real16", "Real32Trunc", "Real32Quant"),
[](const auto &testParam) { return testParam.param; });

TEST(RNTupleMerger, MissingQuantizedFieldOutsideZeroRange)
{
FileRaii first("test_ntuple_merge_zero_quant_first.root");
FileRaii middle("test_ntuple_merge_zero_quant_middle.root");
FileRaii last("test_ntuple_merge_zero_quant_last.root");
FileRaii output("test_ntuple_merge_zero_quant_output.root");
const std::array<const FileRaii *, 3> files{&first, &middle, &last};
for (unsigned int source = 0; source < files.size(); ++source) {
auto model = RNTupleModel::Create();
model->MakeField<int>("entry");
if (source != 1) {
auto field = std::make_unique<RField<float>>("value");
field->SetQuantized(20, {1., 64.});
model->AddField(std::move(field));
}
auto writer = RNTupleWriter::Recreate(std::move(model), "ntuple", files[source]->GetPath());
if (source != 1)
*writer->GetModel().GetDefaultEntry().GetPtr<float>("value") = 10.;
writer->Fill();
}
std::vector<std::unique_ptr<RPageSource>> sources;
std::vector<RPageSource *> sourcePtrs;
for (const auto *file : files) {
sources.push_back(RPageSource::Create("ntuple", file->GetPath()));
sourcePtrs.push_back(sources.back().get());
}
auto destination = std::make_unique<RPageSinkFile>("ntuple", output.GetPath(), RNTupleWriteOptions());
RNTupleMerger merger{std::move(destination)};
RNTupleMergeOptions options;
options.fMergingMode = ENTupleMergingMode::kUnion;
auto result = merger.Merge(sourcePtrs, options);
ASSERT_FALSE(bool(result));
EXPECT_THAT(result.GetError()->GetReport(), testing::HasSubstr("cannot generate default values for field `value`"));
EXPECT_THAT(result.GetError()->GetReport(), testing::HasSubstr("quantization"));
}

struct RNTupleMergerDeferredZeroPages : public ::testing::TestWithParam<std::tuple<std::string, bool>> {};

TEST_P(RNTupleMergerDeferredZeroPages, SecondSource)
{
const auto [type, twoClusters] = GetParam();
auto check = [&](auto sample) {
using T = decltype(sample);
auto setValue = [](T &value, unsigned int entry) {
if constexpr (std::is_same_v<T, std::array<float, 3>>)
value = {float(entry + 1), float(entry + 2), float(entry + 3)};
else
value = T(entry + 1);
};
FileRaii first("test_ntuple_merge_deferred_zero_first.root");
FileRaii second("test_ntuple_merge_deferred_zero_second.root");
FileRaii output("test_ntuple_merge_deferred_zero_output.root");
constexpr unsigned int nEntriesPerSource = 10;
constexpr unsigned int firstDeferredEntry = 5;
{
auto model = RNTupleModel::Create();
model->MakeField<int>("entry");
auto value = model->MakeField<T>("value");
auto writer = RNTupleWriter::Recreate(std::move(model), "ntuple", first.GetPath());
for (unsigned int i = 0; i < nEntriesPerSource; ++i) {
setValue(*value, i);
writer->Fill();
}
}
{
auto model = RNTupleModel::Create();
model->MakeField<int>("entry");
auto writer = RNTupleWriter::Recreate(std::move(model), "ntuple", second.GetPath());
for (unsigned int i = 0; i < firstDeferredEntry; ++i)
writer->Fill();
if (twoClusters)
writer->CommitCluster();
auto updater = writer->CreateModelUpdater();
updater->BeginUpdate();
updater->AddField(std::make_unique<RField<T>>("value"));
updater->CommitUpdate();
auto value = writer->GetModel().GetDefaultEntry().template GetPtr<T>("value");
for (unsigned int i = firstDeferredEntry; i < nEntriesPerSource; ++i) {
setValue(*value, nEntriesPerSource + i);
writer->Fill();
}
}

auto firstSource = RPageSource::Create("ntuple", first.GetPath());
auto secondSource = RPageSource::Create("ntuple", second.GetPath());
std::vector<RPageSource *> sources{firstSource.get(), secondSource.get()};
auto destination = std::make_unique<RPageSinkFile>("ntuple", output.GetPath(), RNTupleWriteOptions());
RNTupleMerger merger{std::move(destination)};
RNTupleMergeOptions options;
options.fMergingMode = ENTupleMergingMode::kUnion;
auto result = merger.Merge(sources, options);
ASSERT_TRUE(bool(result)) << result.GetError()->GetReport();

auto reader = RNTupleReader::Open("ntuple", output.GetPath());
ASSERT_EQ(2 * nEntriesPerSource, reader->GetNEntries());
const auto &descriptor = reader->GetDescriptor();
const auto fieldId = descriptor.FindFieldId(type == "Array3" ? "value._0" : "value");
const auto &field = descriptor.GetFieldDescriptor(fieldId);
ASSERT_EQ(1u, field.GetLogicalColumnIds().size());
const auto columnId = field.GetLogicalColumnIds()[0];
const auto repetitions = type == "Array3" ? 3u : (type == "Bitset8" ? 8u : 1u);
ASSERT_EQ(twoClusters ? 3u : 2u, descriptor.GetNActiveClusters());
for (const auto &cluster : descriptor.GetActiveClusterIterable()) {
ASSERT_EQ(cluster.GetNEntries() * repetitions, cluster.GetColumnRange(columnId).GetNElements());
}
auto value = reader->GetModel().GetDefaultEntry().template GetPtr<T>("value");
for (auto i : reader->GetEntryRange()) {
SCOPED_TRACE(i);
T expected{};
if (i < nEntriesPerSource || i >= nEntriesPerSource + firstDeferredEntry)
setValue(expected, i);
reader->LoadEntry(i);
EXPECT_EQ(expected, *value);
}
};
if (type == "Array3")
check(std::array<float, 3>{});
else if (type == "Bitset8")
check(std::bitset<8>{});
else
check(bool{});
}

INSTANTIATE_TEST_SUITE_P(Repeated, RNTupleMergerDeferredZeroPages,
::testing::Combine(::testing::Values("Array3", "Bitset8", "Bool"), ::testing::Bool()),
[](const auto &testParam) {
return std::get<0>(testParam.param) +
(std::get<1>(testParam.param) ? "TwoClusters" : "OneCluster");
});

struct RNTupleMergerDeferred : public ::testing::TestWithParam<std::tuple<int, int>> {};

TEST_P(RNTupleMergerDeferred, MergeSecondDeferred)
Expand Down