[ntuple] Add RPageStorage::Close

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py

@@ -141,6 +141,12 @@ def _RNTupleReader_LoadEntry(self, *args):
     return self._LoadEntry(*args)
 
 
+def _RNTupleReader_exit(self, *args):
+    import ROOT
+    ROOT.Internal.CloseRNTupleReader(self)
+    return False
+
+
 @pythonization("RNTupleReader", ns="ROOT")
 def pythonize_RNTupleReader(klass):
     klass._Open = klass.Open
@@ -149,6 +155,9 @@ def pythonize_RNTupleReader(klass):
     klass._LoadEntry = klass.LoadEntry
     klass.LoadEntry = _RNTupleReader_LoadEntry
 
+    klass.__enter__ = lambda reader: reader
+    klass.__exit__ = _RNTupleReader_exit
+
 
 def _RNTupleWriter_Append(model, *args):
     # In Python, the user cannot create REntries directly from a model, so we can safely clone it and avoid destructively passing the user argument.
@@ -174,7 +183,9 @@ def _RNTupleWriter_Fill(self, *args):
 
 
 def _RNTupleWriter_exit(self, *args):
+    import ROOT
     self.CommitDataset()
+    ROOT.Internal.CloseRNTupleWriter(self)
     return False
 
 

tree/ntuple/inc/ROOT/RNTupleReader.hxx

@@ -34,6 +34,14 @@
 
 namespace ROOT {
 class RNTuple;
+class RNTupleReader;
+
+namespace Internal {
+
+// Used by the pythonization.
+void CloseRNTupleReader(RNTupleReader &reader);
+
+} // namespace Internal
 
 /// Listing of the different options that can be printed by RNTupleReader::GetInfo()
 enum class ENTupleInfo {
@@ -65,6 +73,8 @@ std::cout << "myNTuple has " << reader->GetNEntries() << " entries\n";
 // clang-format on
 class RNTupleReader {
 private:
+   friend void Internal::CloseRNTupleReader(RNTupleReader &reader);
+
    /// Set as the page source's scheduler for parallel page decompression if implicit multi-threading (IMT) is on.
    /// Needs to be destructed after the page source is destructed (and thus be declared before)
    std::unique_ptr<Internal::RPageStorage::RTaskScheduler> fUnzipTasks;

tree/ntuple/inc/ROOT/RNTupleWriter.hxx

@@ -41,6 +41,9 @@ namespace Internal {
 // Non-public factory method for an RNTuple writer that uses an already constructed page sink
 std::unique_ptr<RNTupleWriter>
 CreateRNTupleWriter(std::unique_ptr<ROOT::RNTupleModel> model, std::unique_ptr<Internal::RPageSink> sink);
+
+// Used by the pythonization.
+void CloseRNTupleWriter(RNTupleWriter &writer);
 } // namespace Internal
 
 // clang-format off
@@ -102,6 +105,7 @@ class RNTupleWriter {
    friend ROOT::RNTupleModel::RUpdater;
    friend std::unique_ptr<RNTupleWriter>
       Internal::CreateRNTupleWriter(std::unique_ptr<ROOT::RNTupleModel>, std::unique_ptr<Internal::RPageSink>);
+   friend void Internal::CloseRNTupleWriter(RNTupleWriter &writer);
 
 private:
    RNTupleFillContext fFillContext;
@@ -117,6 +121,9 @@ private:
    // Helper function that is called from CommitCluster() when necessary
    void CommitClusterGroup();
 
+   // Only called by Python
+   void Close();
+
    /// Create a writer, potentially wrapping the sink in a RPageSinkBuf.
    static std::unique_ptr<RNTupleWriter> Create(std::unique_ptr<ROOT::RNTupleModel> model,
                                                 std::unique_ptr<Internal::RPageSink> sink,

tree/ntuple/inc/ROOT/RPageStorage.hxx

@@ -191,6 +191,9 @@ public:
    /// overriding this.
    virtual ROOT::Experimental::Detail::RNTupleMetrics &GetMetrics() { return fMetrics; }
 
+   // Called only by the Pythonization. Closes the source/sink without destroying it.
+   virtual void Close() {}
+
    /// Returns the NTuple name.
    const std::string &GetNTupleName() const { return fNTupleName; }
 

tree/ntuple/inc/ROOT/RPageStorageDaos.hxx

@@ -141,6 +141,8 @@ protected:
 public:
    RPageSinkDaos(std::string_view ntupleName, std::string_view uri, const ROOT::RNTupleWriteOptions &options);
    ~RPageSinkDaos() override;
+
+   void Close() final;
 }; // class RPageSinkDaos
 
 // clang-format off
@@ -188,6 +190,7 @@ public:
    std::string GetObjectClass() const;
 
    void LoadStreamerInfo() final;
+   void Close() final;
 }; // class RPageSourceDaos
 
 } // namespace Internal

tree/ntuple/inc/ROOT/RPageStorageFile.hxx

@@ -104,6 +104,8 @@ public:
    ~RPageSinkFile() override;
 
    void UpdateSchema(const ROOT::Internal::RNTupleModelChangeset &changeset, ROOT::NTupleSize_t firstEntry) final;
+
+   void Close() final;
 }; // class RPageSinkFile
 
 // clang-format off
@@ -186,6 +188,8 @@ public:
    LoadClusters(std::span<ROOT::Internal::RCluster::RKey> clusterKeys) final;
 
    void LoadStreamerInfo() final;
+
+   void Close() final;
 }; // class RPageSourceFile
 
 } // namespace Internal

tree/ntuple/src/RNTupleReader.cxx

@@ -270,3 +270,8 @@ ROOT::DescriptorId_t ROOT::RNTupleReader::RetrieveFieldId(std::string_view field
    }
    return fieldId;
 }
+
+void ROOT::Internal::CloseRNTupleReader(RNTupleReader &reader)
+{
+   reader.fSource->Close();
+}

tree/ntuple/src/RNTupleWriter.cxx

@@ -150,3 +150,12 @@ ROOT::Internal::CreateRNTupleWriter(std::unique_ptr<ROOT::RNTupleModel> model,
 {
    return std::unique_ptr<ROOT::RNTupleWriter>(new ROOT::RNTupleWriter(std::move(model), std::move(sink)));
 }
+void ROOT::Internal::CloseRNTupleWriter(RNTupleWriter &writer)
+{
+   writer.Close();
+}
+
+void ROOT::RNTupleWriter::Close()
+{
+   fFillContext.fSink->Close();
+}

tree/ntuple/src/RPageStorageDaos.cxx

@@ -483,6 +483,11 @@ void ROOT::Experimental::Internal::RPageSinkDaos::WriteNTupleAnchor()
       kDistributionKeyDefault, kAttributeKeyAnchor, kCidMetadata);
 }
 
+void ROOT::Experimental::Internal::RPageSinkDaos::Close()
+{
+   fDaosContainer.reset();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 ROOT::Experimental::Internal::RPageSourceDaos::RPageSourceDaos(std::string_view ntupleName, std::string_view uri,
@@ -773,3 +778,8 @@ void ROOT::Experimental::Internal::RPageSourceDaos::LoadStreamerInfo()
 {
    R__LOG_WARNING(ROOT::Internal::NTupleLog()) << "DAOS-backed sources have no associated StreamerInfo to load.";
 }
+
+void ROOT::Experimental::Internal::RPageSourceDaos::Close()
+{
+   fDaosContainer.reset();
+}

tree/ntuple/src/RPageStorageFile.cxx

@@ -295,6 +295,11 @@ void ROOT::Internal::RPageSinkFile::CommitDatasetImpl(unsigned char *serializedF
    fWriter->Commit(GetWriteOptions().GetCompression());
 }
 
+void ROOT::Internal::RPageSinkFile::Close()
+{
+   fWriter.reset();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 ROOT::Internal::RPageSourceFile::RPageSourceFile(std::string_view ntupleName, const ROOT::RNTupleReadOptions &opts)
@@ -722,3 +727,8 @@ void ROOT::Internal::RPageSourceFile::LoadStreamerInfo()
 {
    fReader.LoadStreamerInfo();
 }
+
+void ROOT::Internal::RPageSourceFile::Close()
+{
+   fFile.reset();
+}