[ntuple] Merger: avoid resealing if not needed

Currently we either "fast merge" the page (i.e. we copy it as-is to the
destination) or we reseal it if fast merging is not possible.
There is a middle ground where we can avoid resealing if only the page
compression differs but not the on-disk encoding.
We were not currently leveraging this and doing the extra work of
resealing the pages in all cases where just recompression could be done.

With this PR we introduce this middle ground case ("L2 merging", aka
"recompress merging") which skips the resealing (thus the re-encoding)
of the pages whenever possible.

Author: silverweed

tree/ntuple/src/RNTupleMerger.cxx

@@ -267,6 +267,38 @@ struct RChangeCompressionFunc {
    RPageStorage::RSealedPage &fSealedPage;
    ROOT::Internal::RPageAllocator &fPageAlloc;
    std::uint8_t *fBuffer;
+   std::size_t fBufSize;
+
+   void operator()() const
+   {
+      assert(fSrcColElement.GetIdentifier() == fDstColElement.GetIdentifier());
+
+      const auto bytesPacked = fSrcColElement.GetPackedSize(fSealedPage.GetNElements());
+      // TODO: this buffer could be kept and reused across pages
+      auto unzipBuf = MakeUninitArray<unsigned char>(bytesPacked);
+      ROOT::Internal::RNTupleDecompressor::Unzip(fSealedPage.GetBuffer(), fSealedPage.GetDataSize(), bytesPacked,
+                                                 unzipBuf.get());
+
+      // TODO: maybe we should check writeOpts.fWriteChecksum rather than fSealedPage.GetHasChecksum()...
+      const auto checksumSize = fSealedPage.GetHasChecksum() * sizeof(std::uint64_t);
+      assert(fBufSize >= bytesPacked + checksumSize);
+      auto nBytesZipped = ROOT::Internal::RNTupleCompressor::Zip(unzipBuf.get(), bytesPacked,
+                                                                 fMergeOptions.fCompressionSettings.value(), fBuffer);
+
+      fSealedPage = {fBuffer, nBytesZipped + checksumSize, fSealedPage.GetNElements(), fSealedPage.GetHasChecksum()};
+      fSealedPage.ChecksumIfEnabled();
+   }
+};
+
+struct RResealFunc {
+   const RColumnElementBase &fSrcColElement;
+   const RColumnElementBase &fDstColElement;
+   const RNTupleMergeOptions &fMergeOptions;
+
+   RPageStorage::RSealedPage &fSealedPage;
+   ROOT::Internal::RPageAllocator &fPageAlloc;
+   std::uint8_t *fBuffer;
+   std::size_t fBufSize;
 
    void operator()() const
    {
@@ -277,11 +309,25 @@ struct RChangeCompressionFunc {
       sealConf.fBuffer = fBuffer;
       sealConf.fCompressionSettings = *fMergeOptions.fCompressionSettings;
       sealConf.fWriteChecksum = fSealedPage.GetHasChecksum();
+      assert(fBufSize >= fSealedPage.GetDataSize() + fSealedPage.GetHasChecksum() * sizeof(std::uint64_t));
       auto refSealedPage = RPageSink::SealPage(sealConf);
       fSealedPage = refSealedPage;
    }
 };
 
+struct RTaskVisitor {
+   std::optional<ROOT::Experimental::TTaskGroup> &fGroup;
+
+   template <typename T>
+   void operator()(T &f)
+   {
+      if (fGroup)
+         fGroup->Run(f);
+      else
+         f();
+   }
+};
+
 struct RCommonField {
    const ROOT::RFieldDescriptor *fSrc;
    const ROOT::RFieldDescriptor *fDst;
@@ -778,15 +824,23 @@ RNTupleMerger::MergeCommonColumns(ROOT::Internal::RClusterPool &clusterPool,
 
       // Each column range potentially has a distinct compression settings
       const auto colRangeCompressionSettings = clusterDesc.GetColumnRange(columnId).GetCompressionSettings().value();
-      // If either the compression or the encoding of the source doesn't match that of the destination, we need
-      // to reseal the page. Otherwise, if both match, we can fast merge.
+
+      // Select "merging level". There are 3 levels, from fastest to slowest, depending on the case:
+      // L1: compression and encoding of src and dest both match: we can simply copy the page
+      // L2: compression of dest doesn't match the src but encoding does: we must recompress the page but can avoid
+      //     resealing it.
+      // L3: on-disk encoding doesn't match: we need to reseal the page, which implies decompressing and recompressing
+      //     it.
+      const bool compressionIsDifferent =
+         colRangeCompressionSettings != mergeData.fMergeOpts.fCompressionSettings.value();
       const bool needsResealing =
-         colRangeCompressionSettings != mergeData.fMergeOpts.fCompressionSettings.value() ||
          srcColElement->GetIdentifier().fOnDiskType != dstColElement->GetIdentifier().fOnDiskType;
+      const bool needsRecompressing = compressionIsDifferent || needsResealing;
 
-      if (needsResealing && mergeData.fMergeOpts.fExtraVerbose) {
+      if (needsRecompressing && mergeData.fMergeOpts.fExtraVerbose) {
          R__LOG_INFO(NTupleMergeLog())
-            << "Resealing column " << column.fColumnName << ": { compression: " << colRangeCompressionSettings << " => "
+            << (needsResealing ? "Resealing" : "Recompressing") << " column " << column.fColumnName
+            << ": { compression: " << colRangeCompressionSettings << " => "
             << mergeData.fMergeOpts.fCompressionSettings.value()
             << ", onDiskType: " << RColumnElementBase::GetColumnTypeName(srcColElement->GetIdentifier().fOnDiskType)
             << " => " << RColumnElementBase::GetColumnTypeName(dstColElement->GetIdentifier().fOnDiskType);
@@ -795,7 +849,7 @@ RNTupleMerger::MergeCommonColumns(ROOT::Internal::RClusterPool &clusterPool,
       size_t pageBufferBaseIdx = sealedPageData.fBuffers.size();
       // If the column range already has the right compression we don't need to allocate any new buffer, so we don't
       // bother reserving memory for them.
-      if (needsResealing)
+      if (needsRecompressing)
          sealedPageData.fBuffers.resize(sealedPageData.fBuffers.size() + pages.GetPageInfos().size());
 
       // If this column is deferred, we may need to fill "holes" until its real start. We fill any missing entry
@@ -831,18 +885,40 @@ RNTupleMerger::MergeCommonColumns(ROOT::Internal::RClusterPool &clusterPool,
          sealedPage.VerifyChecksumIfEnabled().ThrowOnError();
          R__ASSERT(onDiskPage && (onDiskPage->GetSize() == sealedPage.GetBufferSize()));
 
-         if (needsResealing) {
+         if (needsRecompressing) {
+            std::optional<std::variant<RChangeCompressionFunc, RResealFunc>> task;
             const auto uncompressedSize = srcColElement->GetSize() * sealedPage.GetNElements();
             auto &buffer = sealedPageData.fBuffers[pageBufferBaseIdx + pageIdx];
-            buffer = MakeUninitArray<std::uint8_t>(uncompressedSize + checksumSize);
-            RChangeCompressionFunc compressTask{
-               *srcColElement, *dstColElement, mergeData.fMergeOpts, sealedPage, *fPageAlloc, buffer.get(),
-            };
-
-            if (fTaskGroup)
-               fTaskGroup->Run(compressTask);
-            else
-               compressTask();
+            const auto bufSize = uncompressedSize + checksumSize;
+            // NOTE: we currently allocate the max possible size for this buffer and don't shrink it afterward.
+            // We might want to introduce an option that trades speed for memory usage and shrink the buffer to fit
+            // the actual data size after recompressing.
+            buffer = MakeUninitArray<std::uint8_t>(bufSize);
+
+            if (needsResealing) {
+               task.emplace(RResealFunc{
+                  *srcColElement,
+                  *dstColElement,
+                  mergeData.fMergeOpts,
+                  sealedPage,
+                  *fPageAlloc,
+                  buffer.get(),
+                  bufSize,
+               });
+            } else {
+               task.emplace(RChangeCompressionFunc{
+                  *srcColElement,
+                  *dstColElement,
+                  mergeData.fMergeOpts,
+                  sealedPage,
+                  *fPageAlloc,
+                  buffer.get(),
+                  bufSize,
+               });
+            }
+
+            assert(task.has_value());
+            std::visit(RTaskVisitor{fTaskGroup}, *task);
          }
 
          ++pageIdx;