DPL: allow determining the origin from user provide input

In order to support embedding, we need allow the user to provide
a mapping between the desired origin and the level in the parent file
chain where the table should be found.

Author: ktf

Framework/AnalysisSupport/src/AODJAlienReaderHelpers.cxx

@@ -10,7 +10,10 @@
 // or submit itself to any jurisdiction.
 
 #include "AODJAlienReaderHelpers.h"
+#include <charconv>
 #include <memory>
+#include <ranges>
+#include <vector>
 #include "Framework/TableTreeHelpers.h"
 #include "Framework/AnalysisHelpers.h"
 #include "Framework/DataProcessingStats.h"
@@ -111,10 +114,31 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
   if (ctx.options().isSet("aod-parent-access-level")) {
     parentAccessLevel = ctx.options().get<int>("aod-parent-access-level");
   }
-  auto callback = AlgorithmSpec{adaptStateful([parentFileReplacement, parentAccessLevel](ConfigParamRegistry const& options,
-                                                                                         DeviceSpec const& spec,
-                                                                                         Monitoring& monitoring,
-                                                                                         DataProcessingStats& stats) {
+  std::vector<std::pair<std::string, int>> originLevelMapping;
+  if (ctx.options().isSet("aod-origin-level-mapping")) {
+    auto originLevelMappingStr = ctx.options().get<std::string>("aod-origin-level-mapping");
+    for (auto pairRange : originLevelMappingStr | std::views::split(',')) {
+      std::string_view pair{pairRange.begin(), pairRange.end()};
+      auto colonPos = pair.find(':');
+      if (colonPos == std::string_view::npos) {
+        LOGP(fatal, "Badly formatted aod-origin-level-mapping entry: \"{}\"", pair);
+        continue;
+      }
+      std::string key(pair.substr(0, colonPos));
+      std::string_view valueStr = pair.substr(colonPos + 1);
+      int value{};
+      auto [ptr, ec] = std::from_chars(valueStr.data(), valueStr.data() + valueStr.size(), value);
+      if (ec == std::errc{}) {
+        originLevelMapping.emplace_back(std::move(key), value);
+      } else {
+        LOGP(fatal, "Unable to parse level in aod-origin-level-mapping entry: \"{}\"", pair);
+      }
+    }
+  }
+  auto callback = AlgorithmSpec{adaptStateful([parentFileReplacement, parentAccessLevel, originLevelMapping](ConfigParamRegistry const& options,
+                                                                                                             DeviceSpec const& spec,
+                                                                                                             Monitoring& monitoring,
+                                                                                                             DataProcessingStats& stats) {
     // FIXME: not actually needed, since data processing stats can specify that we should
     // send the initial value.
     stats.updateStats({static_cast<short>(ProcessingStatsId::ARROW_BYTES_CREATED), DataProcessingStats::Op::Set, 0});
@@ -134,7 +158,7 @@ AlgorithmSpec AODJAlienReaderHelpers::rootFileReaderCallback(ConfigContext const
     auto maxRate = options.get<float>("aod-max-io-rate");
 
     // create a DataInputDirector
-    auto didir = std::make_shared<DataInputDirector>(std::vector<std::string>{filename}, DataInputDirectorContext{&monitoring, parentAccessLevel, parentFileReplacement});
+    auto didir = std::make_shared<DataInputDirector>(std::vector<std::string>{filename}, DataInputDirectorContext{&monitoring, parentAccessLevel, parentFileReplacement, originLevelMapping});
     if (options.isSet("aod-reader-json")) {
       auto jsonFile = options.get<std::string>("aod-reader-json");
       if (!didir->readJson(jsonFile)) {

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -122,7 +122,7 @@ void DataInputDescriptor::addFileNameHolder(FileNameHolder* fn)
   mfilenames.emplace_back(fn);
 }
 
-bool DataInputDescriptor::setFile(int counter, std::string_view origin)
+bool DataInputDescriptor::setFile(int counter, int wantedParentLevel, std::string_view origin)
 {
   // no files left
   if (counter >= getNumberInputfiles()) {
@@ -133,7 +133,9 @@ bool DataInputDescriptor::setFile(int counter, std::string_view origin)
   // of the filename. In the future we might expand this for proper rewriting of the
   // filename based on the origin and the original file information.
   std::string filename = mfilenames[counter]->fileName;
-  if (!origin.starts_with("AOD")) {
+  // In case we do not need to remap parent levels, the requested origin is what
+  // drives the filename.
+  if (wantedParentLevel == -1 && !origin.starts_with("AOD")) {
     filename = std::regex_replace(filename, std::regex("[.]root$"), fmt::format("_{}.root", origin));
   }
 
@@ -218,11 +220,11 @@ bool DataInputDescriptor::setFile(int counter, std::string_view origin)
   return true;
 }
 
-uint64_t DataInputDescriptor::getTimeFrameNumber(int counter, int numTF, std::string_view origin)
+uint64_t DataInputDescriptor::getTimeFrameNumber(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin)
 {
 
   // open file
-  if (!setFile(counter, origin)) {
+  if (!setFile(counter, wantedParentLevel, wantedOrigin)) {
     return 0ul;
   }
 
@@ -234,10 +236,27 @@ uint64_t DataInputDescriptor::getTimeFrameNumber(int counter, int numTF, std::st
   return (mfilenames[counter]->listOfTimeFrameNumbers)[numTF];
 }
 
-arrow::dataset::FileSource DataInputDescriptor::getFileFolder(int counter, int numTF, std::string_view origin)
+arrow::dataset::FileSource DataInputDescriptor::getFileFolder(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin)
 {
+  // If mapped to a parent level deeper than current, skip directly to the right level.
+  if (wantedParentLevel != -1 && mLevel < wantedParentLevel) {
+    if (!setFile(counter, wantedParentLevel, wantedOrigin)) {
+      return {};
+    }
+    auto folderName = fmt::format("DF_{}", mfilenames[counter]->listOfTimeFrameNumbers[numTF]);
+    auto parentFile = getParentFile(counter, numTF, "", wantedParentLevel, wantedOrigin);
+    if (parentFile == nullptr) {
+      return {};
+    }
+    int parentNumTF = parentFile->findDFNumber(0, folderName);
+    if (parentNumTF == -1) {
+      return {};
+    }
+    return parentFile->getFileFolder(0, parentNumTF, wantedParentLevel, wantedOrigin);
+  }
+
   // open file
-  if (!setFile(counter, origin)) {
+  if (!setFile(counter, wantedParentLevel, wantedOrigin)) {
     return {};
   }
 
@@ -251,7 +270,7 @@ arrow::dataset::FileSource DataInputDescriptor::getFileFolder(int counter, int n
   return {fmt::format("DF_{}", mfilenames[counter]->listOfTimeFrameNumbers[numTF]), mCurrentFilesystem};
 }
 
-DataInputDescriptor* DataInputDescriptor::getParentFile(int counter, int numTF, std::string treename, std::string_view origin)
+DataInputDescriptor* DataInputDescriptor::getParentFile(int counter, int numTF, std::string treename, int wantedParentLevel, std::string_view wantedOrigin)
 {
   if (!mParentFileMap) {
     // This file has no parent map
@@ -288,7 +307,7 @@ DataInputDescriptor* DataInputDescriptor::getParentFile(int counter, int numTF,
   mParentFile->mdefaultFilenamesPtr = new std::vector<FileNameHolder*>;
   mParentFile->mdefaultFilenamesPtr->emplace_back(makeFileNameHolder(parentFileName->GetString().Data()));
   mParentFile->fillInputfiles();
-  mParentFile->setFile(0, origin);
+  mParentFile->setFile(0, wantedParentLevel, wantedOrigin);
   return mParentFile;
 }
 
@@ -446,8 +465,38 @@ struct CalculateDelta {
 bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh, int counter, int numTF, std::string treename, size_t& totalSizeCompressed, size_t& totalSizeUncompressed)
 {
   CalculateDelta t(mIOTime);
-  std::string origin = dh.dataOrigin.as<std::string>();
-  auto folder = getFileFolder(counter, numTF, origin);
+  std::string wantedOrigin = dh.dataOrigin.as<std::string>();
+  int wantedLevel = -1;
+  for (auto& [origin, level] : mContext.parentLevelToOrigin) {
+    if (origin == wantedOrigin) {
+      wantedLevel = level;
+      break;
+    }
+  }
+
+  // If this origin is mapped to a parent level deeper than current, skip directly without
+  // attempting to read from this level.
+  if (wantedLevel != -1 && mLevel < wantedLevel) {
+    if (!setFile(counter, wantedLevel, wantedOrigin)) {
+      t.deactivate();
+      return false;
+    }
+    auto folderName = fmt::format("DF_{}", mfilenames[counter]->listOfTimeFrameNumbers[numTF]);
+    auto parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
+    if (parentFile == nullptr) {
+      auto rootFS = std::dynamic_pointer_cast<TFileFileSystem>(mCurrentFilesystem);
+      throw std::runtime_error(fmt::format(R"(No parent file found for "{}" while looking for level {} in "{}")", treename, wantedLevel, rootFS->GetFile()->GetName()));
+    }
+    int parentNumTF = parentFile->findDFNumber(0, folderName);
+    if (parentNumTF == -1) {
+      auto parentRootFS = std::dynamic_pointer_cast<TFileFileSystem>(parentFile->mCurrentFilesystem);
+      throw std::runtime_error(fmt::format(R"(DF {} listed in parent file map but not found in the corresponding file "{}")", folderName, parentRootFS->GetFile()->GetName()));
+    }
+    t.deactivate();
+    return parentFile->readTree(outputs, dh, 0, parentNumTF, treename, totalSizeCompressed, totalSizeUncompressed);
+  }
+
+  auto folder = getFileFolder(counter, numTF, wantedLevel, wantedOrigin);
   if (!folder.filesystem()) {
     t.deactivate();
     return false;
@@ -480,7 +529,7 @@ bool DataInputDescriptor::readTree(DataAllocator& outputs, header::DataHeader dh
   if (!format) {
     t.deactivate();
     LOGP(debug, "Could not find tree {}. Trying in parent file.", fullpath.path());
-    auto parentFile = getParentFile(counter, numTF, treename, origin);
+    auto parentFile = getParentFile(counter, numTF, treename, wantedLevel, wantedOrigin);
     if (parentFile != nullptr) {
       int parentNumTF = parentFile->findDFNumber(0, folder.path());
       if (parentNumTF == -1) {
@@ -813,8 +862,15 @@ arrow::dataset::FileSource DataInputDirector::getFileFolder(header::DataHeader d
     didesc = mdefaultDataInputDescriptor;
   }
   std::string origin = dh.dataOrigin.as<std::string>();
+  int wantedLevel = -1;
+  for (auto& [o, level] : mContext.parentLevelToOrigin) {
+    if (o == origin) {
+      wantedLevel = level;
+      break;
+    }
+  }
 
-  return didesc->getFileFolder(counter, numTF, origin);
+  return didesc->getFileFolder(counter, numTF, wantedLevel, origin);
 }
 
 int DataInputDirector::getTimeFramesInFile(header::DataHeader dh, int counter)
@@ -836,8 +892,15 @@ uint64_t DataInputDirector::getTimeFrameNumber(header::DataHeader dh, int counte
     didesc = mdefaultDataInputDescriptor;
   }
   std::string origin = dh.dataOrigin.as<std::string>();
+  int wantedLevel = -1;
+  for (auto& [o, level] : mContext.parentLevelToOrigin) {
+    if (o == origin) {
+      wantedLevel = level;
+      break;
+    }
+  }
 
-  return didesc->getTimeFrameNumber(counter, numTF, origin);
+  return didesc->getTimeFrameNumber(counter, numTF, wantedLevel, origin);
 }
 
 bool DataInputDirector::readTree(DataAllocator& outputs, header::DataHeader dh, int counter, int numTF, size_t& totalSizeCompressed, size_t& totalSizeUncompressed)

Framework/AnalysisSupport/src/DataInputDirector.h

@@ -21,6 +21,7 @@
 #include <arrow/dataset/dataset.h>
 
 #include <regex>
+#include <vector>
 #include "rapidjson/fwd.h"
 
 namespace o2::monitoring
@@ -44,6 +45,7 @@ struct DataInputDirectorContext {
   o2::monitoring::Monitoring* monitoring = nullptr;
   int allowedParentLevel = 0;
   std::string parentFileReplacement = "";
+  std::vector<std::pair<std::string, int>> parentLevelToOrigin = {};
 };
 
 class DataInputDescriptor
@@ -71,7 +73,7 @@ class DataInputDescriptor
 
   void addFileNameHolder(FileNameHolder* fn);
   int fillInputfiles();
-  bool setFile(int counter, std::string_view origin);
+  bool setFile(int counter, int wantedParentLevel, std::string_view wantedOrigin);
 
   // getters
   std::string getInputfilesFilename();
@@ -81,9 +83,9 @@ class DataInputDescriptor
   int getNumberTimeFrames() { return mtotalNumberTimeFrames; }
   int findDFNumber(int file, std::string dfName);
 
-  uint64_t getTimeFrameNumber(int counter, int numTF, std::string_view origin);
-  arrow::dataset::FileSource getFileFolder(int counter, int numTF, std::string_view origin);
-  DataInputDescriptor* getParentFile(int counter, int numTF, std::string treename, std::string_view origin);
+  uint64_t getTimeFrameNumber(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin);
+  arrow::dataset::FileSource getFileFolder(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin);
+  DataInputDescriptor* getParentFile(int counter, int numTF, std::string treename, int wantedParentLevel, std::string_view wantedOrigin);
   int getTimeFramesInFile(int counter);
   int getReadTimeFramesInFile(int counter);
 

Framework/Core/src/Plugin.cxx

@@ -168,6 +168,9 @@ struct DiscoverAODOptionsInCommandLine : o2::framework::ConfigDiscoveryPlugin {
           if (key == "aod-parent-access-level") {
             results.push_back(ConfigParamSpec{"aod-parent-access-level", VariantType::String, value, {"Allow parent file access up to specified level. Default: no (0)"}});
           }
+          if (key == "aod-origin-level-mapping") {
+            results.push_back(ConfigParamSpec{"aod-origin-level-mapping", VariantType::String, value, {"Map origin to parent level for AOD reading. Syntax: ORIGIN:LEVEL[,ORIGIN2:LEVEL2,...]. E.g. \"DYN:1\"."}});
+          }
         }
         if (injectOption) {
           results.push_back(ConfigParamSpec{"aod-writer-compression", VariantType::Int, 505, {"AOD Compression options"}});