@@ -53,7 +53,7 @@ class TTreeFileSystem : public VirtualRootFileSystemBase
5353 const std::string& path,
5454 const std::shared_ptr<const arrow::KeyValueMetadata>& metadata) override ;
5555
56- virtual TTree* GetTree (arrow::dataset::FileSource source) = 0;
56+ virtual std::unique_ptr< TTree>& GetTree (arrow::dataset::FileSource source) = 0;
5757};
5858
5959class SingleTreeFileSystem : public TTreeFileSystem
@@ -72,14 +72,14 @@ class SingleTreeFileSystem : public TTreeFileSystem
7272 return " ttree" ;
7373 }
7474
75- TTree* GetTree (arrow::dataset::FileSource) override
75+ std::unique_ptr< TTree>& GetTree (arrow::dataset::FileSource) override
7676 {
7777 // Simply return the only TTree we have
7878 return mTree ;
7979 }
8080
8181 private:
82- TTree* mTree ;
82+ std::unique_ptr< TTree> mTree ;
8383};
8484
8585arrow::Result<arrow::fs::FileInfo> SingleTreeFileSystem::GetFileInfo (std::string const & path)
@@ -158,7 +158,9 @@ class TTreeFileFormat : public arrow::dataset::FileFormat
158158class TTreeOutputStream : public arrow ::io::OutputStream
159159{
160160 public:
161- TTreeOutputStream (TTree*, std::string branchPrefix);
161+ // Using a pointer means that the tree itself is owned by another
162+ // class
163+ TTreeOutputStream (TTree *, std::string branchPrefix);
162164
163165 arrow::Status Close () override ;
164166
@@ -265,7 +267,7 @@ arrow::Result<arrow::RecordBatchGenerator> TTreeFileFormat::ScanBatchesAsync(
265267 auto fs = std::dynamic_pointer_cast<TTreeFileSystem>(containerFS->GetSubFilesystem (treeFragment->source ()));
266268
267269 int64_t rows = -1 ;
268- TTree* tree = fs->GetTree (treeFragment->source ());
270+ auto & tree = fs->GetTree (treeFragment->source ());
269271 for (auto & field : fields) {
270272 // The field actually on disk
271273 auto physicalField = physical_schema->GetFieldByName (field->name ());
@@ -477,9 +479,9 @@ arrow::Result<std::shared_ptr<arrow::io::OutputStream>> TTreeFileSystem::OpenOut
477479 arrow::dataset::FileSource source{path, shared_from_this ()};
478480 auto prefix = metadata->Get (" branch_prefix" );
479481 if (prefix.ok ()) {
480- return std::make_shared<TTreeOutputStream>(GetTree (source), *prefix);
482+ return std::make_shared<TTreeOutputStream>(GetTree (source). get () , *prefix);
481483 }
482- return std::make_shared<TTreeOutputStream>(GetTree (source), " " );
484+ return std::make_shared<TTreeOutputStream>(GetTree (source). get () , " " );
483485}
484486
485487namespace
@@ -541,7 +543,7 @@ arrow::Result<std::shared_ptr<arrow::Schema>> TTreeFileFormat::Inspect(const arr
541543 if (!treeFs.get ()) {
542544 throw runtime_error_f (" Unknown filesystem %s\n " , source.filesystem ()->type_name ().c_str ());
543545 }
544- TTree* tree = treeFs->GetTree (source);
546+ auto & tree = treeFs->GetTree (source);
545547
546548 auto branches = tree->GetListOfBranches ();
547549 auto n = branches->GetEntries ();
@@ -688,7 +690,7 @@ class TTreeFileWriter : public arrow::dataset::FileWriter
688690 // We already have a tree stream, let's derive a new one
689691 // with the destination_locator_.path as prefix for the branches
690692 // This way we can multiplex multiple tables in the same tree.
691- auto tree = treeStream->GetTree ();
693+ auto * tree = treeStream->GetTree ();
692694 treeStream = std::make_shared<TTreeOutputStream>(tree, destination_locator_.path );
693695 } else {
694696 // I could simply set a prefix here to merge to an already existing tree.
@@ -834,7 +836,7 @@ class TTreeFileWriter : public arrow::dataset::FileWriter
834836 arrow::Future<> FinishInternal () override
835837 {
836838 auto treeStream = std::dynamic_pointer_cast<TTreeOutputStream>(destination_);
837- TTree * tree = treeStream->GetTree ();
839+ auto * tree = treeStream->GetTree ();
838840 tree->Write (" " , TObject::kOverwrite );
839841 tree->SetDirectory (nullptr );
840842
0 commit comments