diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py index b9d3ab06028aa..d62b960285b69 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py +++ b/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 diff --git a/tree/ntuple/inc/ROOT/RNTupleReader.hxx b/tree/ntuple/inc/ROOT/RNTupleReader.hxx index f0aafef634e48..864b453d001df 100644 --- a/tree/ntuple/inc/ROOT/RNTupleReader.hxx +++ b/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 fUnzipTasks; diff --git a/tree/ntuple/inc/ROOT/RNTupleWriter.hxx b/tree/ntuple/inc/ROOT/RNTupleWriter.hxx index 7662a922d1b1e..d646bde41ecf3 100644 --- a/tree/ntuple/inc/ROOT/RNTupleWriter.hxx +++ b/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 CreateRNTupleWriter(std::unique_ptr model, std::unique_ptr 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 Internal::CreateRNTupleWriter(std::unique_ptr, std::unique_ptr); + 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 Create(std::unique_ptr model, std::unique_ptr sink, diff --git a/tree/ntuple/inc/ROOT/RPageStorage.hxx b/tree/ntuple/inc/ROOT/RPageStorage.hxx index c396a3f6ae551..8bfaf0a0cefd8 100644 --- a/tree/ntuple/inc/ROOT/RPageStorage.hxx +++ b/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; } diff --git a/tree/ntuple/inc/ROOT/RPageStorageDaos.hxx b/tree/ntuple/inc/ROOT/RPageStorageDaos.hxx index 4d6200fb5f590..e1f3ad0af9316 100644 --- a/tree/ntuple/inc/ROOT/RPageStorageDaos.hxx +++ b/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 diff --git a/tree/ntuple/inc/ROOT/RPageStorageFile.hxx b/tree/ntuple/inc/ROOT/RPageStorageFile.hxx index bbef8c11b0efc..0f99082d55006 100644 --- a/tree/ntuple/inc/ROOT/RPageStorageFile.hxx +++ b/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 clusterKeys) final; void LoadStreamerInfo() final; + + void Close() final; }; // class RPageSourceFile } // namespace Internal diff --git a/tree/ntuple/src/RNTupleReader.cxx b/tree/ntuple/src/RNTupleReader.cxx index 28b0fb94f1316..4c8baed244535 100644 --- a/tree/ntuple/src/RNTupleReader.cxx +++ b/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(); +} diff --git a/tree/ntuple/src/RNTupleWriter.cxx b/tree/ntuple/src/RNTupleWriter.cxx index 21c21939c915f..4f30a8879c6eb 100644 --- a/tree/ntuple/src/RNTupleWriter.cxx +++ b/tree/ntuple/src/RNTupleWriter.cxx @@ -150,3 +150,12 @@ ROOT::Internal::CreateRNTupleWriter(std::unique_ptr model, { return std::unique_ptr(new ROOT::RNTupleWriter(std::move(model), std::move(sink))); } +void ROOT::Internal::CloseRNTupleWriter(RNTupleWriter &writer) +{ + writer.Close(); +} + +void ROOT::RNTupleWriter::Close() +{ + fFillContext.fSink->Close(); +} diff --git a/tree/ntuple/src/RPageStorageDaos.cxx b/tree/ntuple/src/RPageStorageDaos.cxx index 1de8981b54746..1661749b4834a 100644 --- a/tree/ntuple/src/RPageStorageDaos.cxx +++ b/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(); +} diff --git a/tree/ntuple/src/RPageStorageFile.cxx b/tree/ntuple/src/RPageStorageFile.cxx index dd60247c1658a..2ad1efd389d67 100644 --- a/tree/ntuple/src/RPageStorageFile.cxx +++ b/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(); +}