diff --git a/tree/ntuple/src/RNTupleMerger.cxx b/tree/ntuple/src/RNTupleMerger.cxx index b44247301b4af..29ef952afa226 100644 --- a/tree/ntuple/src/RNTupleMerger.cxx +++ b/tree/ntuple/src/RNTupleMerger.cxx @@ -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(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(parent.GetNRepetitions(), 1); + parentId = parent.GetParentId(); + } + return nElements; +} + // Generates default (zero) values for the given columns ROOT::RResult RNTupleMergeStrategyDefault::GenerateZeroPagesForColumns(size_t nEntriesToGenerate, std::span columns, @@ -1153,21 +1168,8 @@ ROOT::RResult 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(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(parent.GetNRepetitions(), 1); - parentId = parent.GetParentId(); - } - if (skipColumn) + const auto nElementsPerEntry = GetElementsPerEntry(*field, *column.fParentNTupleDescriptor); + if (!nElementsPerEntry) continue; const auto structure = field->GetStructure(); @@ -1184,35 +1186,40 @@ ROOT::RResult 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(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::Success(); @@ -1287,11 +1294,15 @@ ROOT::RResult 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 diff --git a/tree/ntuple/test/ntuple_merger.cxx b/tree/ntuple/test/ntuple_merger.cxx index 76912f77d9346..3aafbd8c6a420 100644 --- a/tree/ntuple/test/ntuple_merger.cxx +++ b/tree/ntuple/test/ntuple_merger.cxx @@ -20,7 +20,9 @@ #include #include +#include #include +#include #include using ROOT::TestSupport::CheckDiagsRAII; @@ -2454,6 +2456,223 @@ TEST(RNTupleMerger, SkipMissing) EXPECT_TRUE(ok); } +struct RNTupleMergerZeroPages : public ::testing::TestWithParam {}; + +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 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("entry"); + if (source != 1) { + auto field = std::make_unique>("value"); + if constexpr (std::is_same_v) { + 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 value; + if (source != 1) + value = writer->GetModel().GetDefaultEntry().template GetPtr("value"); + for (unsigned int i = 0; i < nEntriesPerSource; ++i) { + *entry = source * nEntriesPerSource + i; + if (value) + *value = T(*entry + 1); + writer->Fill(); + } + } + + std::vector> sources; + std::vector 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("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(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("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) + 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 files{&first, &middle, &last}; + for (unsigned int source = 0; source < files.size(); ++source) { + auto model = RNTupleModel::Create(); + model->MakeField("entry"); + if (source != 1) { + auto field = std::make_unique>("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("value") = 10.; + writer->Fill(); + } + std::vector> sources; + std::vector 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("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> {}; + +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>) + 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("entry"); + auto value = model->MakeField("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("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>("value")); + updater->CommitUpdate(); + auto value = writer->GetModel().GetDefaultEntry().template GetPtr("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 sources{firstSource.get(), secondSource.get()}; + auto destination = std::make_unique("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("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{}); + 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> {}; TEST_P(RNTupleMergerDeferred, MergeSecondDeferred)