include/podio/ArrowReader.h

0001 #ifndef PODIO_ARROWREADER_H
0002 #define PODIO_ARROWREADER_H
0003
0004 #include "podio/podioVersion.h"
0005 #include "podio/utilities/ArrowFrameData.h"
0006 #include "podio/utilities/ReaderCommon.h"
0007
0008 #include <cstddef>
0009 #include <filesystem>
0010 #include <map>
0011 #include <memory>
0012 #include <optional>
0013 #include <string>
0014 #include <vector>
0015
0016 namespace arrow {
0017 class Table;
0018 }
0019
0020 namespace podio {
0021
0022 /// Arrow backend reader for PODIO
0023 ///
0024 /// Reads data from a directory structure containing one Parquet file per category
0025 /// and a metadata.json file containing metadata for reading.
0026 class ArrowReader : public ReaderCommon {
0027 public:
0028   /// Create an ArrowReader
0029   ArrowReader();
0030
0031   /// Open the passed directory for reading.
0032   ///
0033   /// @param directory The path to the directory to read from
0034   void openFile(const std::string& directory);
0035
0036   ~ArrowReader() = default;
0037
0038   ArrowReader(const ArrowReader&) = delete;
0039   ArrowReader& operator=(const ArrowReader&) = delete;
0040   ArrowReader(ArrowReader&&) = delete;
0041   ArrowReader& operator=(ArrowReader&&) = delete;
0042
0043   /// Read the next entry for the given category
0044   std::unique_ptr<podio::ArrowFrameData> readNextEntry(std::string_view name,
0045                                                        const std::vector<std::string>& collsToRead = {});
0046
0047   /// Read the specific entry for the given category
0048   std::unique_ptr<podio::ArrowFrameData> readEntry(std::string_view name, size_t index,
0049                                                    const std::vector<std::string>& collsToRead = {});
0050
0051   /// Get the number of entries for a category
0052   size_t getEntries(std::string_view name) const;
0053
0054 private:
0055   struct CategoryInfo {
0056     std::string filePath{};
0057     size_t entries = 0;
0058     size_t currentIndex = 0;
0059     std::shared_ptr<arrow::Table> table{nullptr};
0060   };
0061
0062   void loadCategoryTable(CategoryInfo& catInfo);
0063
0064   std::string m_directory{};
0065   std::map<std::string, CategoryInfo> m_categories{};
0066 };
0067
0068 } // namespace podio
0069
0070 #endif // PODIO_ARROWREADER_H