Skip to content

Commit

Permalink
DPL: write support for TTree using arrow::Dataset API
Browse files Browse the repository at this point in the history
  • Loading branch information
ktf committed Nov 20, 2024
1 parent 15dcbf6 commit 897cc11
Show file tree
Hide file tree
Showing 4 changed files with 441 additions and 3 deletions.
32 changes: 32 additions & 0 deletions Framework/Core/include/Framework/RootArrowFilesystem.h
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,11 @@ class TTreeFileSystem : public VirtualRootFileSystemBase
{
return std::dynamic_pointer_cast<VirtualRootFileSystemBase>(shared_from_this());
};

arrow::Result<std::shared_ptr<arrow::io::OutputStream>> OpenOutputStream(
const std::string& path,
const std::shared_ptr<const arrow::KeyValueMetadata>& metadata) override;

virtual TTree* GetTree(arrow::dataset::FileSource source) = 0;
};

Expand Down Expand Up @@ -128,6 +133,10 @@ class TFileFileSystem : public VirtualRootFileSystemBase

std::shared_ptr<VirtualRootFileSystemBase> GetSubFilesystem(arrow::dataset::FileSource source) override;

arrow::Result<std::shared_ptr<arrow::io::OutputStream>> OpenOutputStream(
const std::string& path,
const std::shared_ptr<const arrow::KeyValueMetadata>& metadata) override;

// We can go back to the TFile in case this is needed.
TDirectoryFile* GetFile()
{
Expand Down Expand Up @@ -218,6 +227,29 @@ class TTreeFileFormat : public arrow::dataset::FileFormat
const std::shared_ptr<arrow::dataset::FileFragment>& fragment) const override;
};

// An arrow outputstream which allows to write to a ttree
class TTreeOutputStream : public arrow::io::OutputStream
{
public:
TTreeOutputStream(TTree* t);

arrow::Status Close() override;

arrow::Result<int64_t> Tell() const override;

arrow::Status Write(const void* data, int64_t nbytes) override;

bool closed() const override;

TTree* GetTree()
{
return mTree;
}

private:
TTree* mTree;
};

} // namespace o2::framework

#endif // O2_FRAMEWORK_ROOT_ARROW_FILESYSTEM_H_
Loading

0 comments on commit 897cc11

Please sign in to comment.