DPL: Move DataInputDirector to arrow::Dataset API

Author: ktf

Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx

@@ -10,10 +10,12 @@
 // or submit itself to any jurisdiction.
 
 #include "AODJAlienReaderHelpers.h"
+#include <memory>
 #include "Framework/TableTreeHelpers.h"
 #include "Framework/AnalysisHelpers.h"
 #include "Framework/DataProcessingStats.h"
 #include "Framework/RootTableBuilderHelpers.h"
+#include "Framework/RootArrowFilesystem.h"
 #include "Framework/AlgorithmSpec.h"
 #include "Framework/ConfigParamRegistry.h"
 #include "Framework/ControlService.h"
@@ -41,6 +43,8 @@
 #include <arrow/io/interfaces.h>
 #include <arrow/table.h>
 #include <arrow/util/key_value_metadata.h>
+#include <arrow/dataset/dataset.h>
+#include <arrow/dataset/file_base.h>
 
 using namespace o2;
 using namespace o2::aod;
@@ -272,11 +276,13 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
             // Origin file name for derived output map
             auto o2 = Output(TFFileNameHeader);
             auto fileAndFolder = didir->getFileFolder(dh, fcnt, ntf);
-            std::string currentFilename(fileAndFolder.file->GetName());
-            if (strcmp(fileAndFolder.file->GetEndpointUrl()->GetProtocol(), "file") == 0 && fileAndFolder.file->GetEndpointUrl()->GetFile()[0] != '/') {
+            auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(fileAndFolder.filesystem());
+            auto* f = dynamic_cast<TFile*>(rootFS->GetFile());
+            std::string currentFilename(f->GetFile()->GetName());
+            if (strcmp(f->GetEndpointUrl()->GetProtocol(), "file") == 0 && f->GetEndpointUrl()->GetFile()[0] != '/') {
               // This is not an absolute local path. Make it absolute.
               static std::string pwd = gSystem->pwd() + std::string("/");
-              currentFilename = pwd + std::string(fileAndFolder.file->GetName());
+              currentFilename = pwd + std::string(f->GetName());
             }
             outputs.make<std::string>(o2) = currentFilename;
           }
@@ -312,7 +318,9 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
       auto concrete = DataSpecUtils::asConcreteDataMatcher(firstRoute.matcher);
       auto dh = header::DataHeader(concrete.description, concrete.origin, concrete.subSpec);
       auto fileAndFolder = didir->getFileFolder(dh, fcnt, ntf);
-      if (!fileAndFolder.file) {
+
+      // In case the filesource is empty, move to the next one.
+      if (fileAndFolder.path().empty()) {
         fcnt += 1;
         ntf = 0;
         if (didir->atEnd(fcnt)) {

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -15,6 +15,10 @@
 
 #include "Framework/DataDescriptorMatcher.h"
 #include "Framework/DataAllocator.h"
+#include "Framework/RootArrowFilesystem.h"
+
+#include <arrow/filesystem/filesystem.h>
+#include <arrow/dataset/dataset.h>
 
 #include <regex>
 #include "rapidjson/fwd.h"
@@ -31,16 +35,10 @@ struct FileNameHolder {
   std::string fileName;
   int numberOfTimeFrames = 0;
   std::vector<uint64_t> listOfTimeFrameNumbers;
-  std::vector<std::string> listOfTimeFrameKeys;
   std::vector<bool> alreadyRead;
 };
 FileNameHolder* makeFileNameHolder(std::string fileName);
 
-struct FileAndFolder {
-  TFile* file = nullptr;
-  std::string folderName = "";
-};
-
 class DataInputDescriptor
 {
   /// Holds information concerning the reading of an aod table.
@@ -52,7 +50,6 @@ class DataInputDescriptor
   std::string treename = "";
   std::unique_ptr<data_matcher::DataDescriptorMatcher> matcher;
 
-  DataInputDescriptor() = default;
   DataInputDescriptor(bool alienSupport, int level, o2::monitoring::Monitoring* monitoring = nullptr, int allowedParentLevel = 0, std::string parentFileReplacement = "");
 
   void printOut();
@@ -78,7 +75,7 @@ class DataInputDescriptor
   int findDFNumber(int file, std::string dfName);
 
   uint64_t getTimeFrameNumber(int counter, int numTF);
-  FileAndFolder getFileFolder(int counter, int numTF);
+  arrow::dataset::FileSource getFileFolder(int counter, int numTF);
   DataInputDescriptor* getParentFile(int counter, int numTF, std::string treename);
   int getTimeFramesInFile(int counter);
   int getReadTimeFramesInFile(int counter);
@@ -90,6 +87,7 @@ class DataInputDescriptor
   bool isAlienSupportOn() { return mAlienSupport; }
 
  private:
+  o2::framework::RootObjectReadingFactory mFactory;
   std::string minputfilesFile = "";
   std::string* minputfilesFilePtr = nullptr;
   std::string mFilenameRegex = "";
@@ -98,7 +96,7 @@ class DataInputDescriptor
   std::string mParentFileReplacement;
   std::vector<FileNameHolder*> mfilenames;
   std::vector<FileNameHolder*>* mdefaultFilenamesPtr = nullptr;
-  TFile* mcurrentFile = nullptr;
+  std::shared_ptr<arrow::fs::FileSystem> mCurrentFilesystem;
   int mCurrentFileID = -1;
   bool mAlienSupport = false;
 
@@ -127,7 +125,6 @@ class DataInputDirector
   ~DataInputDirector();
 
   void reset();
-  void createDefaultDataInputDescriptor();
   void printOut();
   bool atEnd(int counter);
 
@@ -140,10 +137,11 @@ class DataInputDirector
   // getters
   DataInputDescriptor* getDataInputDescriptor(header::DataHeader dh);
   int getNumberInputDescriptors() { return mdataInputDescriptors.size(); }
+  void createDefaultDataInputDescriptor();
 
   bool readTree(DataAllocator& outputs, header::DataHeader dh, int counter, int numTF, size_t& totalSizeCompressed, size_t& totalSizeUncompressed);
   uint64_t getTimeFrameNumber(header::DataHeader dh, int counter, int numTF);
-  FileAndFolder getFileFolder(header::DataHeader dh, int counter, int numTF);
+  arrow::dataset::FileSource getFileFolder(header::DataHeader dh, int counter, int numTF);
   int getTimeFramesInFile(header::DataHeader dh, int counter);
 
   uint64_t getTotalSizeCompressed();

Framework/AnalysisSupport/src/DataInputDirector.h

@@ -13,6 +13,7 @@
 #include "Framework/Plugins.h"
 #include "Framework/Signpost.h"
 #include "Framework/Endian.h"
+#include <arrow/buffer.h>
 #include <arrow/dataset/file_base.h>
 #include <arrow/extension_type.h>
 #include <arrow/status.h>
@@ -286,6 +287,8 @@ arrow::Result<arrow::RecordBatchGenerator> TTreeFileFormat::ScanBatchesAsync(
 
   auto generator = [pool = options->pool, treeFragment, dataset_schema, &totalCompressedSize = mTotCompressedSize,
                     &totalUncompressedSize = mTotUncompressedSize]() -> arrow::Future<std::shared_ptr<arrow::RecordBatch>> {
+    O2_SIGNPOST_ID_FROM_POINTER(tid, root_arrow_fs, treeFragment->GetTree());
+    O2_SIGNPOST_START(root_arrow_fs, tid, "Generator", "Creating batch for tree %{public}s", treeFragment->GetTree()->GetName());
     std::vector<std::shared_ptr<arrow::Array>> columns;
     std::vector<std::shared_ptr<arrow::Field>> fields = dataset_schema->fields();
     auto physical_schema = *treeFragment->ReadPhysicalSchema();
@@ -299,27 +302,48 @@ arrow::Result<arrow::RecordBatchGenerator> TTreeFileFormat::ScanBatchesAsync(
 
     for (int fi = 0; fi < dataset_schema->num_fields(); ++fi) {
       auto dataset_field = dataset_schema->field(fi);
+      // This is needed because for now the dataset_field
+      // is actually the schema of the ttree
+      O2_SIGNPOST_EVENT_EMIT(root_arrow_fs, tid, "Generator", "Processing dataset field %{public}s.", dataset_field->name().c_str());
       int physicalFieldIdx = physical_schema->GetFieldIndex(dataset_field->name());
 
       if (physicalFieldIdx < 0) {
         throw runtime_error_f("Cannot find physical field associated to %s", dataset_field->name().c_str());
       }
       if (physicalFieldIdx > 1 && physical_schema->field(physicalFieldIdx - 1)->name().ends_with("_size")) {
+        O2_SIGNPOST_EVENT_EMIT(root_arrow_fs, tid, "Generator", "Field %{public}s has sizes in %{public}s.", dataset_field->name().c_str(),
+                               physical_schema->field(physicalFieldIdx - 1)->name().c_str());
         mappings.push_back({physicalFieldIdx, physicalFieldIdx - 1, fi});
       } else {
+        if (physicalFieldIdx > 1) {
+          O2_SIGNPOST_EVENT_EMIT(root_arrow_fs, tid, "Generator", "Field %{public}s previous field is %{public}s.", dataset_field->name().c_str(),
+                                 physical_schema->field(physicalFieldIdx - 1)->name().c_str());
+        }
         mappings.push_back({physicalFieldIdx, -1, fi});
       }
     }
 
     auto* tree = treeFragment->GetTree();
-    tree->SetCacheSize(25000000);
     auto branches = tree->GetListOfBranches();
+    size_t totalTreeSize = 0;
+    std::vector<TBranch*> selectedBranches;
     for (auto& mapping : mappings) {
-      tree->AddBranchToCache((TBranch*)branches->At(mapping.mainBranchIdx), false);
+      selectedBranches.push_back((TBranch*)branches->At(mapping.mainBranchIdx));
+      O2_SIGNPOST_EVENT_EMIT(root_arrow_fs, tid, "Generator", "Adding branch %{public}s to stream.", selectedBranches.back()->GetName());
+      totalTreeSize += selectedBranches.back()->GetTotalSize();
       if (mapping.vlaIdx != -1) {
-        tree->AddBranchToCache((TBranch*)branches->At(mapping.vlaIdx), false);
+        selectedBranches.push_back((TBranch*)branches->At(mapping.vlaIdx));
+        O2_SIGNPOST_EVENT_EMIT(root_arrow_fs, tid, "Generator", "Adding branch %{public}s to stream.", selectedBranches.back()->GetName());
+        totalTreeSize += selectedBranches.back()->GetTotalSize();
       }
     }
+
+    size_t cacheSize = std::max(std::min(totalTreeSize, 25000000UL), 1000000UL);
+    O2_SIGNPOST_EVENT_EMIT(root_arrow_fs, tid, "Generator", "Resizing cache to %zu.", cacheSize);
+    tree->SetCacheSize(cacheSize);
+    for (auto* branch : selectedBranches) {
+      tree->AddBranchToCache(branch, false);
+    }
     tree->StopCacheLearningPhase();
 
     static TBufferFile buffer{TBuffer::EMode::kWrite, 4 * 1024 * 1024};
@@ -400,9 +424,7 @@ arrow::Result<arrow::RecordBatchGenerator> TTreeFileFormat::ScanBatchesAsync(
         }
       } else {
         // This is needed for branches which have not been persisted.
-        auto bytes = branch->GetTotBytes();
-        auto branchSize = bytes ? bytes : 1000000;
-        auto&& result = arrow::AllocateResizableBuffer(branchSize, pool);
+        auto&& result = arrow::AllocateBuffer(branch->GetTotalSize(), pool);
         if (!result.ok()) {
           throw runtime_error("Cannot allocate values buffer");
         }
@@ -423,7 +445,7 @@ arrow::Result<arrow::RecordBatchGenerator> TTreeFileFormat::ScanBatchesAsync(
         if (mapping.vlaIdx != -1) {
           auto* mSizeBranch = (TBranch*)branches->At(mapping.vlaIdx);
           offsetBuffer = std::make_unique<TBufferFile>(TBuffer::EMode::kWrite, 4 * 1024 * 1024);
-          result = arrow::AllocateResizableBuffer((totalEntries + 1) * (int64_t)sizeof(int), pool);
+          result = arrow::AllocateBuffer((totalEntries + 1) * (int64_t)sizeof(int), pool);
           if (!result.ok()) {
             throw runtime_error("Cannot allocate offset buffer");
           }
@@ -435,6 +457,9 @@ arrow::Result<arrow::RecordBatchGenerator> TTreeFileFormat::ScanBatchesAsync(
           // read sizes first
           while (readEntries < totalEntries) {
             auto readLast = mSizeBranch->GetBulkRead().GetEntriesSerialized(readEntries, *offsetBuffer);
+            if (readLast == -1) {
+              throw runtime_error_f("Unable to read from branch %s.", mSizeBranch->GetName());
+            }
             readEntries += readLast;
             for (auto i = 0; i < readLast; ++i) {
               offsets[count++] = (int)offset;
@@ -492,6 +517,7 @@ arrow::Result<arrow::RecordBatchGenerator> TTreeFileFormat::ScanBatchesAsync(
     auto batch = arrow::RecordBatch::Make(dataset_schema, rows, columns);
     totalCompressedSize += tree->GetZipBytes();
     totalUncompressedSize += tree->GetTotBytes();
+    O2_SIGNPOST_END(root_arrow_fs, tid, "Generator", "Done creating batch compressed:%zu uncompressed:%zu", totalCompressedSize, totalUncompressedSize);
     return batch;
   };
   return generator;

Framework/AnalysisSupport/src/TTreePlugin.cxx

@@ -144,6 +144,8 @@ class TFileFileSystem : public VirtualRootFileSystemBase
 
   TFileFileSystem(TDirectoryFile* f, size_t readahead, RootObjectReadingFactory&);
 
+  ~TFileFileSystem() override;
+
   std::string type_name() const override
   {
     return "TDirectoryFile";

Framework/Core/include/Framework/RootArrowFilesystem.h

@@ -42,6 +42,12 @@ TFileFileSystem::TFileFileSystem(TDirectoryFile* f, size_t readahead, RootObject
   ((TFile*)mFile)->SetReadaheadSize(50 * 1024 * 1024);
 }
 
+TFileFileSystem::~TFileFileSystem()
+{
+  mFile->Close();
+  delete mFile;
+}
+
 std::shared_ptr<RootObjectHandler> TFileFileSystem::GetObjectHandler(arrow::dataset::FileSource source)
 {
   // We use a plugin to create the actual objects inside the

Framework/Core/src/RootArrowFilesystem.cxx

@@ -40,7 +40,7 @@ struct EtaAndClsHistogramsSimple {
   OutputObj<TH2F> etaClsH{TH2F("eta_vs_pt", "#eta vs pT", 102, -2.01, 2.01, 100, 0, 10)};
   Produces<o2::aod::SkimmedExampleTrack> skimEx;
 
-  void process(aod::Tracks const& tracks)
+  void process(aod::Tracks const& tracks, aod::FT0s const&)
   {
     LOGP(info, "Invoking the simple one");
     for (auto& track : tracks) {
@@ -54,7 +54,7 @@ struct EtaAndClsHistogramsIUSimple {
   OutputObj<TH2F> etaClsH{TH2F("eta_vs_pt", "#eta vs pT", 102, -2.01, 2.01, 100, 0, 10)};
   Produces<o2::aod::SkimmedExampleTrack> skimEx;
 
-  void process(aod::TracksIU const& tracks)
+  void process(aod::TracksIU const& tracks, aod::FT0s const&)
   {
     LOGP(info, "Invoking the simple one");
     for (auto& track : tracks) {