DataInputDirector.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
#ifndef O2_FRAMEWORK_DATAINPUTDIRECTOR_H_
#define O2_FRAMEWORK_DATAINPUTDIRECTOR_H_
 
#include "TFile.h"
 
#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 <vector>
#include "rapidjson/fwd.h"
 
namespace o2::monitoring
{
class Monitoring;
}
 
namespace o2::framework
{
 
struct FileNameHolder {
  std::string fileName;
  int numberOfTimeFrames = 0;
  std::vector<uint64_t> listOfTimeFrameNumbers;
  std::vector<bool> alreadyRead;
};
 
FileNameHolder makeFileNameHolder(std::string fileName);
 
struct DataInputDirectorContext {
  o2::monitoring::Monitoring* monitoring = nullptr;
  int allowedParentLevel = 0;
  std::string parentFileReplacement = "";
  std::vector<std::pair<std::string, int>> parentLevelToOrigin = {};
  // Optional registry of pre-opened TFiles (keyed by name) used to bypass
  // TFile::Open for testing with in-memory TMemFile instances.
  std::vector<std::pair<std::string, TFile*>> openFiles = {};
 
  int levelForOrigin(std::string_view origin) const
  {
    for (auto& [o, level] : parentLevelToOrigin) {
      if (o == origin) {
        return level;
      }
    }
    return -1;
  }
};
 
class DataInputDescriptor
{
 
 public:
  std::string tablename = "";
  std::string treename = "";
  std::shared_ptr<data_matcher::DataDescriptorMatcher> matcher;
 
  DataInputDescriptor(bool alienSupport, int level, DataInputDirectorContext& context);
  DataInputDescriptor(DataInputDescriptor const&) = default;
 
  void printOut() const;
 
  // setters
  void setInputfilesFile(std::string dffn) { minputfilesFile = dffn; }
  void setInputfilesFile(std::string* dffnptr) { minputfilesFilePtr = dffnptr; }
  void setFilenamesRegex(std::string fn) { mFilenameRegex = fn; }
  void setFilenamesRegex(std::string* fnptr) { mFilenameRegexPtr = fnptr; }
 
  void setDefaultInputfiles(std::vector<FileNameHolder> difnptr) { mdefaultFilenamesPtr = difnptr; }
 
  void addFileNameHolder(FileNameHolder fn);
  int fillInputfiles();
  bool setFile(int counter, int wantedParentLevel, std::string_view wantedOrigin);
 
  // getters
  std::string getInputfilesFilename() const;
  std::string getFilenamesRegexString() const;
  std::regex getFilenamesRegex();
  int getNumberInputfiles() { return mfilenames.size(); }
  int getNumberTimeFrames() const { return mtotalNumberTimeFrames; }
  int findDFNumber(int file, std::string dfName);
 
  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);
  // Open the current file to populate the parent map, then return the parent descriptor and
  // the TF index within it that corresponds to numTF at this level. Returns {nullptr, -1} on failure.
  std::pair<std::shared_ptr<DataInputDescriptor>, int> navigateToLevel(int counter, int numTF, int wantedParentLevel, std::string_view wantedOrigin);
  std::shared_ptr<DataInputDescriptor> getParentFile(int counter, int numTF, std::string treename, int wantedParentLevel, std::string_view wantedOrigin);
  int getTimeFramesInFile(int counter);
  int getReadTimeFramesInFile(int counter);
 
  bool readTree(DataAllocator& outputs, header::DataHeader dh, int counter, int numTF, std::string treename, size_t& totalSizeCompressed, size_t& totalSizeUncompressed);
 
  void printFileOpening();
  void printFileStatistics();
  void closeInputFile();
  bool isAlienSupportOn() { return mAlienSupport; }
 
 private:
  o2::framework::RootObjectReadingFactory mFactory;
  std::string minputfilesFile;
  std::string* minputfilesFilePtr = nullptr;
  std::string mFilenameRegex;
  std::string* mFilenameRegexPtr = nullptr;
  std::vector<FileNameHolder> mfilenames;
  std::vector<FileNameHolder> mdefaultFilenamesPtr;
  std::shared_ptr<arrow::fs::FileSystem> mCurrentFilesystem;
  int mCurrentFileID = -1;
  bool mAlienSupport = false;
 
  DataInputDirectorContext& mContext;
  TMap* mParentFileMap = nullptr;
  std::shared_ptr<DataInputDescriptor> mParentFile = nullptr;
  int mLevel = 0; // level of parent files
 
  int mtotalNumberTimeFrames = 0;
 
  uint64_t mIOTime = 0;
  uint64_t mCurrentFileStartedAt = 0;
};
 
class DataInputDirector
{
 
 public:
  DataInputDirector(std::vector<std::string> inputFiles, DataInputDirectorContext&& context);
  ~DataInputDirector();
 
  void reset();
  void printOut();
  bool atEnd(int counter);
 
  // setters
  void setInputfilesFile(std::string iffn) { minputfilesFile = iffn; }
  void setFilenamesRegex(std::string dfn) { mFilenameRegex = dfn; }
  bool readJson(std::string const& fnjson);
  void closeInputFiles();
 
  // 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, bool wasAOD);
  uint64_t getTimeFrameNumber(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();
  uint64_t getTotalSizeUncompressed();
 
  int getLevelForOrigin(header::DataOrigin origin) const;
 
 private:
  DataInputDirectorContext mContext;
  std::string minputfilesFile;
  std::string* const minputfilesFilePtr = &minputfilesFile;
  std::string mFilenameRegex;
  std::string* const mFilenameRegexPtr = &mFilenameRegex;
  std::shared_ptr<DataInputDescriptor> mdefaultDataInputDescriptor = nullptr;
  std::vector<FileNameHolder> mdefaultInputFiles;
  std::vector<DataInputDescriptor> mdataInputDescriptors;
 
  bool mDebugMode = false;
  bool mAlienSupport = false;
 
  bool readJsonDocument(rapidjson::Document* doc);
  bool isValid();
};
 
} // namespace o2::framework
 
#endif // O2_FRAMEWORK_DATAINPUTDIRECTOR_H_