DPL: add support for decompressing directly to shared memory

This PR postpones the read operations which would usually populate an intermediate
RecordBatch and it performs them directly on its subsequent shared memory
serialization. Doing so avoids having the intermediate representation allocate most
of the memory.

For the moment this is only done for TTree. RNtuple support will come in a subsequent PR.

Author: ktf

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -399,6 +399,7 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   }
   // FIXME: Ugly. We should detect the format from the treename, good enough for now.
   std::shared_ptr<arrow::dataset::FileFormat> format;
+  FragmentToBatch::StreamerCreator creator = nullptr;
 
   auto fullpath = arrow::dataset::FileSource{folder.path() + "/" + treename, folder.filesystem()};
 
@@ -407,6 +408,7 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
     void* handle = capability.getHandle(rootFS, objectPath);
     if (handle) {
       format = capability.factory().format();
+      creator = capability.factory().deferredOutputStreamer;
       break;
     }
   }
@@ -449,13 +451,12 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
 
   // FIXME: This should allow me to create a memory pool
   // which I can then use to scan the dataset.
-  //
-  auto f2b = outputs.make<FragmentToBatch>(o);
+  auto f2b = outputs.make<FragmentToBatch>(o, creator, *fragment);
 
   //// add branches to read
   //// fill the table
   f2b->setLabel(treename.c_str());
-  f2b->fill(*fragment, datasetSchema, format);
+  f2b->fill(datasetSchema, format);
 
   mIOTime += (uv_hrtime() - ioStart);
 

Framework/Core/include/Framework/TableTreeHelpers.h

@@ -11,6 +11,8 @@
 #ifndef O2_FRAMEWORK_TABLETREEHELPERS_H_
 #define O2_FRAMEWORK_TABLETREEHELPERS_H_
 
+#include <arrow/buffer.h>
+#include <arrow/io/interfaces.h>
 #include <arrow/record_batch.h>
 #include "TFile.h"
 #include "TTreeReader.h"
@@ -146,15 +148,25 @@ class TreeToTable
 class FragmentToBatch
 {
  public:
-  FragmentToBatch(arrow::MemoryPool* pool = arrow::default_memory_pool());
+  // The function to be used to create the required stream.
+  using StreamerCreator = std::function<std::shared_ptr<arrow::io::OutputStream>(std::shared_ptr<arrow::dataset::FileFragment>, const std::shared_ptr<arrow::ResizableBuffer>& buffer)>;
+
+  FragmentToBatch(StreamerCreator, std::shared_ptr<arrow::dataset::FileFragment>, 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>);
+  void fill(std::shared_ptr<arrow::Schema> dataSetSchema, std::shared_ptr<arrow::dataset::FileFormat>);
   std::shared_ptr<arrow::RecordBatch> finalize();
 
+  std::shared_ptr<arrow::io::OutputStream> streamer(std::shared_ptr<arrow::ResizableBuffer> buffer)
+  {
+    return mCreator(mFragment, buffer);
+  }
+
  private:
+  std::shared_ptr<arrow::dataset::FileFragment> mFragment;
   arrow::MemoryPool* mArrowMemoryPool = nullptr;
   std::string mTableLabel;
   std::shared_ptr<arrow::RecordBatch> mRecordBatch;
+  StreamerCreator mCreator;
 };
 
 // -----------------------------------------------------------------------------

Framework/Core/src/DataAllocator.cxx

@@ -211,34 +211,6 @@ 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>();
@@ -318,16 +290,35 @@ void DataAllocator::adopt(const Output& spec, LifetimeHolder<FragmentToBatch>& f
     // 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());
+    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 = buffer->Reserve(expectedSize);
+    if (reserve.ok() == false) {
+      throw std::runtime_error("Unable to reserve memory for table");
+    }
+
+    auto deferredWriterStream = source.streamer(buffer);
+
+    auto outBatch = arrow::ipc::MakeStreamWriter(deferredWriterStream, 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");
+    }
     // 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.
+    // the LifetimeHolder goes out of scope. See code above.
   };
 
   context.addBuffer(std::move(header), buffer, std::move(finalizer), routeIndex);

Framework/Core/src/TableTreeHelpers.cxx

@@ -13,7 +13,6 @@
 #include "Framework/Endian.h"
 #include "Framework/Signpost.h"
 
-#include "arrow/type_traits.h"
 #include <arrow/dataset/file_base.h>
 #include <arrow/record_batch.h>
 #include <arrow/type.h>
@@ -533,7 +532,7 @@ void TreeToTable::setLabel(const char* label)
   mTableLabel = label;
 }
 
-void TreeToTable::fill(TTree*tree)
+void TreeToTable::fill(TTree* tree)
 {
   std::vector<std::shared_ptr<arrow::ChunkedArray>> columns;
   std::vector<std::shared_ptr<arrow::Field>> fields;
@@ -569,8 +568,10 @@ std::shared_ptr<arrow::Table> TreeToTable::finalize()
   return mTable;
 }
 
-FragmentToBatch::FragmentToBatch(arrow::MemoryPool* pool)
-  : mArrowMemoryPool{pool}
+FragmentToBatch::FragmentToBatch(StreamerCreator creator, std::shared_ptr<arrow::dataset::FileFragment> fragment, arrow::MemoryPool* pool)
+  : mFragment{std::move(fragment)},
+    mArrowMemoryPool{pool},
+    mCreator{std::move(creator)}
 {
 }
 
@@ -579,13 +580,14 @@ 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)
+void FragmentToBatch::fill(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 scanner = format->ScanBatchesAsync(options, mFragment);
   auto batch = (*scanner)();
   mRecordBatch = *batch.result();
+  // Notice that up to here the buffer was not yet filled.
 }
 
 std::shared_ptr<arrow::RecordBatch> FragmentToBatch::finalize()