Skip to content

Commit 7af865a

Browse files
committed
DPL: write support for TTree using arrow::Dataset API
1 parent 15dcbf6 commit 7af865a

File tree

4 files changed

+441
-3
lines changed

4 files changed

+441
-3
lines changed

Framework/Core/include/Framework/RootArrowFilesystem.h

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,11 @@ class TTreeFileSystem : public VirtualRootFileSystemBase
8787
{
8888
return std::dynamic_pointer_cast<VirtualRootFileSystemBase>(shared_from_this());
8989
};
90+
91+
arrow::Result<std::shared_ptr<arrow::io::OutputStream>> OpenOutputStream(
92+
const std::string& path,
93+
const std::shared_ptr<const arrow::KeyValueMetadata>& metadata) override;
94+
9095
virtual TTree* GetTree(arrow::dataset::FileSource source) = 0;
9196
};
9297

@@ -128,6 +133,10 @@ class TFileFileSystem : public VirtualRootFileSystemBase
128133

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

136+
arrow::Result<std::shared_ptr<arrow::io::OutputStream>> OpenOutputStream(
137+
const std::string& path,
138+
const std::shared_ptr<const arrow::KeyValueMetadata>& metadata) override;
139+
131140
// We can go back to the TFile in case this is needed.
132141
TDirectoryFile* GetFile()
133142
{
@@ -218,6 +227,29 @@ class TTreeFileFormat : public arrow::dataset::FileFormat
218227
const std::shared_ptr<arrow::dataset::FileFragment>& fragment) const override;
219228
};
220229

230+
// An arrow outputstream which allows to write to a ttree
231+
class TTreeOutputStream : public arrow::io::OutputStream
232+
{
233+
public:
234+
TTreeOutputStream(TTree* t);
235+
236+
arrow::Status Close() override;
237+
238+
arrow::Result<int64_t> Tell() const override;
239+
240+
arrow::Status Write(const void* data, int64_t nbytes) override;
241+
242+
bool closed() const override;
243+
244+
TTree* GetTree()
245+
{
246+
return mTree;
247+
}
248+
249+
private:
250+
TTree* mTree;
251+
};
252+
221253
} // namespace o2::framework
222254

223255
#endif // O2_FRAMEWORK_ROOT_ARROW_FILESYSTEM_H_

0 commit comments

Comments
 (0)