include/podio/ArrowWriter.h

0001 #ifndef PODIO_ARROWWRITER_H
0002 #define PODIO_ARROWWRITER_H
0003
0004 #include "podio/utilities/DatamodelRegistryIOHelpers.h"
0005
0006 #include <parquet/arrow/writer.h>
0007
0008 #include <cstddef>
0009 #include <cstdint>
0010 #include <filesystem>
0011 #include <map>
0012 #include <memory>
0013 #include <string>
0014 #include <vector>
0015
0016 // Forward declarations for Arrow and Parquet
0017 namespace arrow {
0018 class Schema;
0019 class Table;
0020 } // namespace arrow
0021
0022 namespace podio {
0023
0024 class Frame;
0025
0026 /// Arrow backend writer for PODIO
0027 ///
0028 /// Writes data to a directory structure containing one Parquet file per category
0029 /// and a metadata.json file containing metadata for reading.
0030 class ArrowWriter {
0031 public:
0032   /// Configure the ArrowWriter
0033   struct Options {
0034     size_t maxBufferedRows = 1000;
0035     std::string compression = "";
0036   };
0037
0038   /// Create a ArrowWriter to write to a directory.
0039   ///
0040   /// @note Will create the directory if it doesn't exist. Will throw if it exists
0041   ///       and is not empty.
0042   ///
0043   /// @param directory The path to the output directory.
0044   /// @param options   Configuration options for buffering and compression.
0045   ArrowWriter(const std::string& directory, const Options& options);
0046   explicit ArrowWriter(const std::string& directory);
0047
0048   /// Destructor writes metadata and closes files.
0049   ~ArrowWriter();
0050
0051   ArrowWriter(const ArrowWriter&) = delete;
0052   ArrowWriter& operator=(const ArrowWriter&) = delete;
0053   ArrowWriter(ArrowWriter&&) = delete;
0054   ArrowWriter& operator=(ArrowWriter&&) = delete;
0055
0056   /// Store the given frame with the given category.
0057   ///
0058   /// @note All frames of the same category must have the same collection
0059   ///       names and schemas. Trying to write a frame with different
0060   ///       collections to an existing category will result in an exception.
0061   void writeFrame(const podio::Frame& frame, std::string_view category);
0062
0063   /// Store the given Frame with the given category, specifying collections.
0064   ///
0065   /// @note All frames of the same category must have the same collection
0066   ///       names and schemas. Trying to write a frame with different
0067   ///       collections to an existing category will result in an exception.
0068   void writeFrame(const podio::Frame& frame, std::string_view category, const std::vector<std::string>& collsToWrite);
0069
0070   /// Write the current directory including metadata.json and close files.
0071   void finish();
0072
0073 private:
0074   /// Helper struct to manage category state
0075   struct CategoryInfo {
0076     std::string filePath{};
0077     std::shared_ptr<arrow::Schema> schema{nullptr};
0078     std::vector<std::string> collsToWrite{};
0079     std::vector<std::string> collTypes{};
0080     std::vector<bool> collIsSubset{};
0081     std::vector<uint32_t> collSchemaVersions{};
0082     std::vector<uint32_t> collIDs{};
0083     std::vector<std::shared_ptr<arrow::Table>> buffer{};
0084     std::unique_ptr<parquet::arrow::FileWriter> writer{nullptr};
0085     size_t entries = 0;
0086
0087     CategoryInfo() = default;
0088     ~CategoryInfo() = default;
0089     CategoryInfo(CategoryInfo&&) = default;
0090     CategoryInfo& operator=(CategoryInfo&&) = default;
0091   };
0092
0093   void flushCategory(CategoryInfo& catInfo);
0094   void writeMetadata();
0095   void validateSchema(const CategoryInfo& catInfo, const podio::Frame& frame,
0096                       const std::vector<std::string>& collsToWrite);
0097
0098   std::string m_directory{};
0099   Options m_options{};
0100   std::map<std::string, CategoryInfo> m_categories{};
0101   DatamodelDefinitionCollector m_datamodelCollector{};
0102   bool m_finished = false;
0103 };
0104
0105 } // namespace podio
0106
0107 #endif // PODIO_ARROWWRITER_H