DPL: introduce FragmentToBatch

Author: ktf

Framework/Core/include/Framework/DataAllocator.h

@@ -242,6 +242,15 @@ class DataAllocator
     return t2t;
   }
 
+  template <typename T, typename... Args>
+    requires(requires { static_cast<struct FragmentToBatch>(std::declval<std::decay_t<T>>()); })
+  decltype(auto) make(const Output& spec, Args... args)
+  {
+    auto f2b = std::move(LifetimeHolder<FragmentToBatch>(new std::decay_t<T>(args...)));
+    adopt(spec, f2b);
+    return f2b;
+  }
+
   template <typename T>
     requires is_messageable<T>::value && (!is_specialization_v<T, UninitializedVector>)
   decltype(auto) make(const Output& spec)
@@ -284,6 +293,11 @@ class DataAllocator
   void
     adopt(const Output& spec, LifetimeHolder<struct TreeToTable>&);
 
+  /// Adopt a Source2Batch in the framework and serialise / send
+  /// it as an Arrow Dataset to all consumers of @a spec once done
+  void
+    adopt(const Output& spec, LifetimeHolder<struct FragmentToBatch>&);
+
   /// Adopt an Arrow table and send it to all consumers of @a spec
   void
     adopt(const Output& spec, std::shared_ptr<class arrow::Table>);

Framework/Core/include/Framework/TableTreeHelpers.h

@@ -11,11 +11,14 @@
 #ifndef O2_FRAMEWORK_TABLETREEHELPERS_H_
 #define O2_FRAMEWORK_TABLETREEHELPERS_H_
 
+#include <arrow/record_batch.h>
 #include "TFile.h"
 #include "TTreeReader.h"
 #include "TTreeReaderValue.h"
 #include "TTreeReaderArray.h"
 #include "TableBuilder.h"
+#include <arrow/dataset/file_base.h>
+#include <memory>
 
 // =============================================================================
 namespace o2::framework
@@ -140,6 +143,20 @@ class TreeToTable
   void addReader(TBranch* branch, std::string const& name, bool VLA);
 };
 
+class FragmentToBatch
+{
+ public:
+  FragmentToBatch(arrow::MemoryPool* pool = arrow::default_memory_pool());
+  void setLabel(const char* label);
+  void fill(std::shared_ptr<arrow::dataset::FileFragment>, std::shared_ptr<arrow::Schema> dataSetSchema, std::shared_ptr<arrow::dataset::FileFormat>);
+  std::shared_ptr<arrow::RecordBatch> finalize();
+
+ private:
+  arrow::MemoryPool* mArrowMemoryPool = nullptr;
+  std::string mTableLabel;
+  std::shared_ptr<arrow::RecordBatch> mRecordBatch;
+};
+
 // -----------------------------------------------------------------------------
 } // namespace o2::framework
 

Framework/Core/src/DataAllocator.cxx

@@ -211,6 +211,34 @@ void doWriteTable(std::shared_ptr<FairMQResizableBuffer> b, arrow::Table* table)
   }
 }
 
+void doWriteBatch(std::shared_ptr<FairMQResizableBuffer> b, arrow::RecordBatch* batch)
+{
+  auto mock = std::make_shared<arrow::io::MockOutputStream>();
+  int64_t expectedSize = 0;
+  auto mockWriter = arrow::ipc::MakeStreamWriter(mock.get(), batch->schema());
+  arrow::Status outStatus = mockWriter.ValueOrDie()->WriteRecordBatch(*batch);
+
+  expectedSize = mock->Tell().ValueOrDie();
+  auto reserve = b->Reserve(expectedSize);
+  if (reserve.ok() == false) {
+    throw std::runtime_error("Unable to reserve memory for table");
+  }
+
+  auto stream = std::make_shared<FairMQOutputStream>(b);
+  // This is a copy maybe we can finally get rid of it by having using the
+  // dataset API?
+  auto outBatch = arrow::ipc::MakeStreamWriter(stream.get(), batch->schema());
+  if (outBatch.ok() == false) {
+    throw ::std::runtime_error("Unable to create batch writer");
+  }
+
+  outStatus = outBatch.ValueOrDie()->WriteRecordBatch(*batch);
+
+  if (outStatus.ok() == false) {
+    throw std::runtime_error("Unable to Write batch");
+  }
+}
+
 void DataAllocator::adopt(const Output& spec, LifetimeHolder<TableBuilder>& tb)
 {
   auto& timingInfo = mRegistry.get<TimingInfo>();
@@ -273,6 +301,38 @@ void DataAllocator::adopt(const Output& spec, LifetimeHolder<TreeToTable>& t2t)
   context.addBuffer(std::move(header), buffer, std::move(finalizer), routeIndex);
 }
 
+void DataAllocator::adopt(const Output& spec, LifetimeHolder<FragmentToBatch>& f2b)
+{
+  auto& timingInfo = mRegistry.get<TimingInfo>();
+  RouteIndex routeIndex = matchDataHeader(spec, timingInfo.timeslice);
+
+  auto header = headerMessageFromOutput(spec, routeIndex, o2::header::gSerializationMethodArrow, 0);
+  auto& context = mRegistry.get<ArrowContext>();
+
+  auto creator = [transport = context.proxy().getOutputTransport(routeIndex)](size_t s) -> std::unique_ptr<fair::mq::Message> {
+    return transport->CreateMessage(s);
+  };
+  auto buffer = std::make_shared<FairMQResizableBuffer>(creator);
+
+  f2b.callback = [buffer = buffer, transport = context.proxy().getOutputTransport(routeIndex)](FragmentToBatch& source) {
+    // Serialization happens in here, so that we can
+    // get rid of the intermediate tree 2 table object, saving memory.
+    auto batch = source.finalize();
+    doWriteBatch(buffer, batch.get());
+    // deletion happens in the caller
+  };
+
+  /// To finalise this we write the table to the buffer.
+  /// FIXME: most likely not a great idea. We should probably write to the buffer
+  ///        directly in the TableBuilder, incrementally.
+  auto finalizer = [](std::shared_ptr<FairMQResizableBuffer> b) -> void {
+    // This is empty because we already serialised the object when
+    // the LifetimeHolder goes out of scope.
+  };
+
+  context.addBuffer(std::move(header), buffer, std::move(finalizer), routeIndex);
+}
+
 void DataAllocator::adopt(const Output& spec, std::shared_ptr<arrow::Table> ptr)
 {
   auto& timingInfo = mRegistry.get<TimingInfo>();

Framework/Core/src/TableTreeHelpers.cxx

@@ -13,9 +13,13 @@
 #include "Framework/Endian.h"
 
 #include "arrow/type_traits.h"
+#include <arrow/dataset/file_base.h>
+#include <arrow/record_batch.h>
+#include <arrow/type.h>
 #include <arrow/util/key_value_metadata.h>
 #include <TBufferFile.h>
 
+#include <memory>
 #include <utility>
 namespace TableTreeHelpers
 {
@@ -407,7 +411,7 @@ std::shared_ptr<TTree> TableToTree::process()
 
   for (auto& reader : mColumnReaders) {
     int idealBasketSize = 1024 + reader->fieldSize() * reader->columnEntries(); // minimal additional size needed, otherwise we get 2 baskets
-    int basketSize = std::max(32000, idealBasketSize);        // keep a minimum value
+    int basketSize = std::max(32000, idealBasketSize);                          // keep a minimum value
     // std::cout << "Setting baskets size for " << reader->branchName() << " to " << basketSize << " =  1024 + "
     //           << reader->fieldSize() << " * " << reader->columnEntries() << ". mRows was " << mRows << std::endl;
     mTree->SetBasketSize(reader->branchName(), basketSize);
@@ -555,4 +559,28 @@ std::shared_ptr<arrow::Table> TreeToTable::finalize()
   return mTable;
 }
 
+FragmentToBatch::FragmentToBatch(arrow::MemoryPool* pool)
+  : mArrowMemoryPool{pool}
+{
+}
+
+void FragmentToBatch::setLabel(const char* label)
+{
+  mTableLabel = label;
+}
+
+void FragmentToBatch::fill(std::shared_ptr<arrow::dataset::FileFragment> fragment, std::shared_ptr<arrow::Schema> schema, std::shared_ptr<arrow::dataset::FileFormat> format)
+{
+  auto options = std::make_shared<arrow::dataset::ScanOptions>();
+  options->dataset_schema = schema;
+  auto scanner = format->ScanBatchesAsync(options, fragment);
+  auto batch = (*scanner)();
+  mRecordBatch = *batch.result();
+}
+
+std::shared_ptr<arrow::RecordBatch> FragmentToBatch::finalize()
+{
+  return mRecordBatch;
+}
+
 } // namespace o2::framework