[ntuple] Support merging columns with same type and different metadata

Author: silverweed

tree/ntuple/inc/ROOT/RColumnElementBase.hxx

@@ -17,6 +17,7 @@
 #include <ROOT/RError.hxx>
 #include <ROOT/RFloat16.hxx>
 #include <ROOT/RNTupleTypes.hxx>
+#include <ROOT/RNTupleDescriptor.hxx>
 
 #include <Byteswap.h>
 #include <TError.h>
@@ -147,6 +148,13 @@ std::unique_ptr<RColumnElementBase> RColumnElementBase::Generate(ROOT::ENTupleCo
 template <>
 std::unique_ptr<RColumnElementBase> RColumnElementBase::Generate<void>(ROOT::ENTupleColumnType onDiskType);
 
+struct RColumnFormat {
+   ENTupleColumnType fType = ENTupleColumnType::kUnknown;
+   // 0 means "use default". Only valid for fixed-bitwidth column types.
+   std::uint16_t fBitWidth = 0;
+   std::optional<RColumnDescriptor::RValueRange> fValueRange;
+};
+
 } // namespace ROOT::Internal
 
 #endif

tree/ntuple/inc/ROOT/RPageStorage.hxx

@@ -543,16 +543,10 @@ public:
    [[nodiscard]] std::unique_ptr<RNTupleModel>
    InitFromDescriptor(const ROOT::RNTupleDescriptor &descriptor, bool copyClusters);
 
-   struct RColumnReprElement {
-      ENTupleColumnType fType = ENTupleColumnType::kUnknown;
-      // 0 means "use default". Only valid for fixed-bitwidth column types.
-      std::uint16_t fBitWidth = 0;
-      std::optional<RColumnDescriptor::RValueRange> fValueRange;
-   };
    /// Adds a new column representation to the given field.
    /// \return The physical id of the first newly added column.
-   ROOT::DescriptorId_t
-   AddColumnRepresentation(const ROOT::RFieldDescriptor &field, std::span<const RColumnReprElement> newRepresentation);
+   ROOT::DescriptorId_t AddColumnRepresentation(const ROOT::RFieldDescriptor &field,
+                                                std::span<const ROOT::Internal::RColumnFormat> newRepresentation);
 
    /// Adds a new alias column pointing to an existing column with the given physical id to the given field.
    void AddAliasColumn(const ROOT::RNTupleDescriptor &desc, const ROOT::RFieldDescriptor &field,

tree/ntuple/src/RNTupleMerger.cxx

@@ -345,7 +345,7 @@ struct RColReprMapping {
 /// Note that this also adds a mapping for each new representation, which is why it inherits RColReprMapping.
 struct RColReprExtension : RColReprMapping {
    /// The new representations to be added
-   std::vector<ROOT::Internal::RPagePersistentSink::RColumnReprElement> fSourceRepr;
+   std::vector<ROOT::Internal::RColumnFormat> fSourceRepr;
 };
 
 static std::optional<std::uint32_t>
@@ -500,26 +500,19 @@ static void MatchColumnRepresentations(const ROOT::RNTupleDescriptor &srcDesc, c
                const auto &srcCol = srcDesc.GetColumnDescriptor(srcColId);
                const auto dstColId = dstColumns[dstReprIdx * dstColCardinality + reprColIdx];
                const auto &dstCol = dstDesc.GetColumnDescriptor(dstColId);
-               if (srcCol.GetBitsOnStorage() != dstCol.GetBitsOnStorage() ||
+               if (srcCol.GetType() != dstCol.GetType() ||
+                   srcCol.GetBitsOnStorage() != dstCol.GetBitsOnStorage() ||
                    srcCol.GetValueRange() != dstCol.GetValueRange()) {
-                  std::stringstream ss;
-                  ss << "Source field `" << srcField.GetFieldName()
-                     << "` has a matching column representation as its destination field, however one or "
-                        "more "
-                        "of its columns have different column metadata (bit width and/or value range). "
-                        "Merging variable-sized columns is currently only supported if all metadata is "
-                        "identical between source and destination columns."
-                     << "\n   bit width src: " << srcCol.GetBitsOnStorage() << ", dst: " << dstCol.GetBitsOnStorage()
-                     << ""
-                     << "\n   value range src: " << srcCol.GetValueRange() << ", dst: " << dstCol.GetValueRange();
-                  errors.push_back(ss.str());
+                  matches = false;
                   break;
                }
             }
 
-            // We found a valid matching representation: break the loop on the dst column representations.
-            matchingRepr = dstReprIdx;
-            break;
+            if (matches) {
+               // We found a valid matching representation.
+               matchingRepr = dstReprIdx;
+               break;
+            }
          }
       }
 
@@ -532,7 +525,7 @@ static void MatchColumnRepresentations(const ROOT::RNTupleDescriptor &srcDesc, c
          } else if (matchingRepr < 0) {
             // this representation was not found in the destination
             assert(dstNColReprs < std::numeric_limits<std::uint32_t>::max());
-            std::vector<ROOT::Internal::RPagePersistentSink::RColumnReprElement> newRepr;
+            std::vector<ROOT::Internal::RColumnFormat> newRepr;
             newRepr.reserve(srcColCardinality);
             for (auto reprColIdx = 0u; reprColIdx < srcColCardinality; ++reprColIdx) {
                const auto srcColId = srcColumns[srcReprIdx * srcColCardinality + reprColIdx];

tree/ntuple/src/RPageStorage.cxx

@@ -1146,7 +1146,7 @@ ROOT::Internal::RPagePersistentSink::InitFromDescriptor(const ROOT::RNTupleDescr
 
 ROOT::DescriptorId_t
 ROOT::Internal::RPagePersistentSink::AddColumnRepresentation(const ROOT::RFieldDescriptor &field,
-                                                             std::span<const RColumnReprElement> newRepresentation)
+                                                             std::span<const RColumnFormat> newRepresentation)
 {
    const auto &descriptor = fDescriptorBuilder.GetDescriptor();
 

tree/ntuple/test/ntuple_merger.cxx

@@ -4275,10 +4275,17 @@ TEST(RNTupleMerger, MergeReal32Trunc)
             RNTupleMergeOptions opts;
             opts.fMergingMode = mmode;
             auto res = merger.Merge(sourcePtrs, opts);
-            // Currently we're not supporting merging columns with the same type but different metadata.
-            // TODO: support this.
-            EXPECT_FALSE(bool(res));
-            EXPECT_THAT(res.GetError()->GetReport(), testing::HasSubstr("have different column metadata"));
+            EXPECT_TRUE(bool(res));
+         }
+         {
+            auto reader = ROOT::RNTupleReader::Open("ntuple", fileGuardOut.GetPath());
+            EXPECT_EQ(reader->GetNEntries(), 20);
+            EXPECT_EQ(reader->GetDescriptor().GetNPhysicalColumns(), 2);
+            auto pFlt = reader->GetModel().GetDefaultEntry().GetPtr<float>("flt");
+            for (auto i : reader->GetEntryRange()) {
+               reader->LoadEntry(i);
+               EXPECT_NEAR(*pFlt, i, 0.01f);
+            }
          }
       }
    }
@@ -4333,10 +4340,17 @@ TEST(RNTupleMerger, MergeReal32Quant)
             RNTupleMergeOptions opts;
             opts.fMergingMode = mmode;
             auto res = merger.Merge(sourcePtrs, opts);
-            // Currently we're not supporting merging columns with the same type but different metadata.
-            // TODO: support this.
-            ASSERT_FALSE(bool(res));
-            EXPECT_THAT(res.GetError()->GetReport(), testing::HasSubstr("have different column metadata"));
+            EXPECT_TRUE(bool(res));
+         }
+         {
+            auto reader = ROOT::RNTupleReader::Open("ntuple", fileGuardOut.GetPath());
+            EXPECT_EQ(reader->GetNEntries(), 20);
+            EXPECT_EQ(reader->GetDescriptor().GetNPhysicalColumns(), 2);
+            auto pFlt = reader->GetModel().GetDefaultEntry().GetPtr<float>("flt");
+            for (auto i : reader->GetEntryRange()) {
+               reader->LoadEntry(i);
+               EXPECT_NEAR(*pFlt, i, 0.01f);
+            }
          }
       }
    }
@@ -4406,3 +4420,68 @@ TEST(RNTupleMerger, MergeReal32TruncQuantMixed)
       }
    }
 }
+
+TEST(RNTupleMerger, MergeRealRegularQuantMixed)
+{
+   // Merge two files, both containing the same field, but with the first being a SplitReal64 and the second Real32Quant
+   FileRaii fileGuard1("test_ntuple_merge_realregquant_in_1.root");
+   {
+      auto model = RNTupleModel::Create();
+      auto fieldDbl = model->MakeField<double>("dbl");
+      auto ntuple = RNTupleWriter::Recreate(std::move(model), "ntuple", fileGuard1.GetPath());
+      for (int i = 0; i < 10; ++i) {
+         *fieldDbl = i;
+         ntuple->Fill();
+      }
+   }
+   FileRaii fileGuard2("test_ntuple_merge_realregquant_in_2.root");
+   {
+      auto model = RNTupleModel::Create();
+      auto field = std::make_unique<RField<double>>("dbl");
+      field->SetQuantized(29, {0., 20.});
+      model->AddField(std::move(field));
+      auto ntuple = RNTupleWriter::Recreate(std::move(model), "ntuple", fileGuard2.GetPath());
+      auto fieldDbl = ntuple->GetModel().GetDefaultEntry().GetPtr<double>("dbl");
+      for (int i = 0; i < 10; ++i) {
+         *fieldDbl = 10 + i;
+         ntuple->Fill();
+      }
+   }
+   {
+      // Gather the input sources
+      std::vector<std::unique_ptr<RPageSource>> sources;
+      sources.push_back(RPageSource::Create("ntuple", fileGuard1.GetPath(), RNTupleReadOptions()));
+      sources.push_back(RPageSource::Create("ntuple", fileGuard2.GetPath(), RNTupleReadOptions()));
+      std::vector<RPageSource *> sourcePtrs;
+      for (const auto &s : sources) {
+         sourcePtrs.push_back(s.get());
+      }
+
+      // Now merge the inputs
+      for (const auto mmode : {ENTupleMergingMode::kFilter, ENTupleMergingMode::kStrict, ENTupleMergingMode::kUnion}) {
+         SCOPED_TRACE(std::string("with merging mode = ") + ToString(mmode));
+         FileRaii fileGuardOut("test_ntuple_merge_realregquant_out.root");
+         {
+            auto destination = std::make_unique<RPageSinkFile>("ntuple", fileGuardOut.GetPath(), RNTupleWriteOptions());
+            RNTupleMerger merger{std::move(destination)};
+            RNTupleMergeOptions opts;
+            opts.fMergingMode = mmode;
+            auto res = merger.Merge(sourcePtrs, opts);
+            EXPECT_TRUE(bool(res));
+         }
+         {
+            auto reader = ROOT::RNTupleReader::Open("ntuple", fileGuardOut.GetPath());
+            EXPECT_EQ(reader->GetNEntries(), 20);
+            EXPECT_EQ(reader->GetDescriptor().GetNPhysicalColumns(), 2);
+            auto pDbl = reader->GetModel().GetDefaultEntry().GetPtr<double>("dbl");
+            for (auto i : reader->GetEntryRange()) {
+               reader->LoadEntry(i);
+               if (i < 10)
+                  EXPECT_DOUBLE_EQ(*pDbl, i);
+               else
+                  EXPECT_NEAR(*pDbl, i, 0.01f);
+            }
+         }
+      }
+   }
+}