diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index ec4b9db9b..4e292bbf9 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -29,6 +29,7 @@ jobs: --preset ci-build \ -DENABLE_JULIA=ON \ -DENABLE_ARROW=ON \ + -DENABLE_PARQUET=ON \ -DUSE_EXTERNAL_CATCH2=OFF ln -s build/compile_commands.json compile_commands.json echo "::endgroup::" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 37278ecc5..2f9096474 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -38,6 +38,7 @@ jobs: cmake --preset ci-build \ -DENABLE_JULIA=ON \ -DENABLE_ARROW=ON \ + -DENABLE_PARQUET=$([[ ${{ matrix.LCG }} == dev* ]] && echo "ON" || echo "OFF") \ -DENABLE_RNTUPLE=$([[ ${{ matrix.LCG }} == LCG_104/* ]] && echo "OFF" || echo "ON") \ -DPODIO_RUN_STRACE_TEST=$([[ ${{ matrix.LCG }} == LCG_104/* ]] && echo "OFF" || echo "ON") \ -DCMAKE_INSTALL_PREFIX=$(pwd)/install \ diff --git a/.gitignore b/.gitignore index f239e60f7..222bec2ca 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,8 @@ tests/unittests/Manifest.toml *.root *.dat *.sio +*.parquet +*.podio_arrow # Spack build folders spack* diff --git a/CMakeLists.txt b/CMakeLists.txt index 9b957bbad..d37aae8a4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -68,6 +68,7 @@ ADD_CLANG_TIDY() option(CREATE_DOC "Whether or not to create doxygen doc target." OFF) option(ENABLE_SIO "Build SIO I/O support" OFF) option(ENABLE_ARROW "Build Arrow I/O support" OFF) +option(ENABLE_PARQUET "Build Parquet support" OFF) option(PODIO_RELAX_PYVER "Do not require exact python version match with ROOT" OFF) option(ENABLE_RNTUPLE "Build with support for the new ROOT NTtuple format" OFF) option(ENABLE_DATASOURCE "Build podio's ROOT DataSource" OFF) @@ -147,6 +148,10 @@ if(ENABLE_SIO) endif() endif() +if(ENABLE_PARQUET AND NOT ENABLE_ARROW) + message(FATAL_ERROR "ENABLE_PARQUET requires ENABLE_ARROW to be ON") +endif() + # optionally build with Arrow ----------------------------------------------- if(ENABLE_ARROW) set(lz4Alt_FIND_QUIETLY TRUE) @@ -159,7 +164,19 @@ if(ENABLE_ARROW) message(FATAL_ERROR "Found Arrow, but no usable Arrow CMake target was exported") endif() - message(STATUS "Found Arrow library - will build Arrow I/O support") + if(ENABLE_PARQUET) + find_package(nlohmann_json 3.10 REQUIRED) + find_package(Parquet REQUIRED) + if(TARGET Parquet::parquet_shared) + set(PODIO_PARQUET_TARGET Parquet::parquet_shared) + else() + set(PODIO_PARQUET_TARGET Parquet::parquet_static) + endif() + message(STATUS "Found Arrow and Parquet libraries - will build Arrow I/O support") + else() + message(STATUS "Found Arrow library - will build Arrow I/O support without Parquet") + endif() + list(APPEND PODIO_IO_HANDLERS ARROW) endif() diff --git a/cmake/podioConfig.cmake.in b/cmake/podioConfig.cmake.in index bd487e1dc..674602faa 100644 --- a/cmake/podioConfig.cmake.in +++ b/cmake/podioConfig.cmake.in @@ -47,9 +47,14 @@ if(PODIO_ENABLE_SIO) endif() SET(PODIO_ENABLE_ARROW @ENABLE_ARROW@) +SET(PODIO_ENABLE_PARQUET @ENABLE_PARQUET@) if(PODIO_ENABLE_ARROW) set(lz4Alt_FIND_QUIETLY TRUE) find_dependency(Arrow) + if(PODIO_ENABLE_PARQUET) + find_dependency(nlohmann_json) + find_dependency(Parquet) + endif() set(PODIO_ARROW_TARGET @PODIO_ARROW_TARGET@) endif() diff --git a/cmake/podioMacros.cmake b/cmake/podioMacros.cmake index 7cba9584e..f504f7fce 100644 --- a/cmake/podioMacros.cmake +++ b/cmake/podioMacros.cmake @@ -362,7 +362,7 @@ endif() list(FILTER SOURCES INCLUDE REGEX .*ArrowMapper.cc) add_library(${CORE_LIB}PodioArrow SHARED ${SOURCES}) - target_link_libraries(${CORE_LIB}PodioArrow PUBLIC ${CORE_LIB} podio::podio ${PODIO_ARROW_TARGET}) + target_link_libraries(${CORE_LIB}PodioArrow PUBLIC ${CORE_LIB} podio::podio podio::podioArrow) target_include_directories(${CORE_LIB}PodioArrow PUBLIC $ $) diff --git a/doc/reading_writing.md b/doc/reading_writing.md index 4b32cae67..b45a836ca 100644 --- a/doc/reading_writing.md +++ b/doc/reading_writing.md @@ -1,12 +1,12 @@ -# ROOT Support in Podio +# Reading and writing podio files -Podio supports ROOT file I/O through multiple backends in both C++ and Python -for ROOT TTrees and ROOT RNTuples. Below are the recommended approaches for -reading and writing these files, as well as additional notes to guide usage. +Podio supports file I/O through multiple backends in both C++ and Python. Below +are the recommended approaches for reading and writing these files, as well as +additional notes to guide usage. ## C++ -Podio provides generic and format-specific I/O interfaces for ROOT files. +Podio provides generic and format-specific I/O interfaces. ### Reading @@ -55,8 +55,9 @@ By default, TTrees are written. To explicitly select an output backend, provide the type: ```cpp -auto ttreeWriter = podio::makeWriter(filename, "root"); // Use TTree -auto rntupleWriter = podio::makeWriter(filename, "rntuple"); // Use RNTuple +auto ttreeWriter = podio::makeWriter(filename, "root"); // Use TTree +auto rntupleWriter = podio::makeWriter(filename, "rntuple"); // Use RNTuple +auto parquetWriter = podio::makeWriter(filename, "parquet"); // Use Arrow/Parquet ``` The format can also be set by the environment variable `PODIO_DEFAULT_WRITE_RNTUPLE`. If @@ -66,6 +67,7 @@ the environment variable is set **to a non-empty string**, RNTuples will be the - `.root`: Uses the default backend (TTree or RNTuple if `PODIO_DEFAULT_WRITE_RNTUPLE` is set to a non-empty string), unless specified. - `.sio`: Uses the SIO writer. +- `.podio_parquet`: Uses the Arrow/Parquet writer. - Other extensions are not allowed. Specific writers for each backend are also available: @@ -73,11 +75,50 @@ Specific writers for each backend are also available: ```cpp #include // For TTree output #include // For RNTuple output +#include // For Arrow/Parquet output -podio::ROOTWriter writer(filename); // For TTree output -podio::RNTupleWriter rntupleWriter(filename); // For RNTuple output +podio::ROOTWriter writer(filename); // For TTree output +podio::RNTupleWriter rntupleWriter(filename); // For RNTuple output +podio::ArrowWriter parquetWriter(directoryName); // For Arrow/Parquet output ``` +### Arrow/Parquet I/O + +The Arrow/Parquet backend writes a directory with the `.podio_parquet` +extension. The directory contains one Parquet file per category and a +`metadata.json` file with the podio metadata needed for reading. + +```cpp +#include +#include + +auto writer = podio::makeWriter("events.podio_parquet", "parquet"); +writer.writeFrame(frame, podio::Category::Event); +writer.finish(); + +auto reader = podio::makeReader("events.podio_parquet"); +auto event = reader.readEvent(0); +``` + +The backend-specific classes can also be used directly: + +```cpp +#include +#include + +podio::ArrowWriter writer("events.podio_parquet"); +writer.writeFrame(frame, podio::Category::Event); +writer.finish(); + +podio::ArrowReader reader; +reader.openFile("events.podio_parquet"); +auto event = podio::Frame(reader.readEntry(podio::Category::Event, 0)); +``` + +The default compression for Arrow/Parquet output is controlled at configure +time with the `PODIO_ARROW_DEFAULT_COMPRESSION` CMake option. Supported values +are `UNCOMPRESSED`, `SNAPPY`, and `ZSTD`. + ```{note} Note that the generic readers and writers have methods that are not available in the backend-specific classes. For example, the generic reader has a @@ -111,9 +152,11 @@ Alternatively, instantiate backend-specific readers explicitly: ```python from podio.root_io import Reader # For TTrees from podio.root_io import RNTupleReader # For RNTuples +from podio.arrow_io import Reader as ArrowReader reader = Reader(filename) # For TTree files (.root) rntuple_reader = RNTupleReader(filename) # For RNTuple files (.root) +arrow_reader = ArrowReader(directory) # For Arrow/Parquet directories (.podio_parquet) ``` ### Writing @@ -123,9 +166,12 @@ Similarly, use the appropriate writer class for the file format: ```python from podio.root_io import Writer # For TTrees from podio.root_io import RNTupleWriter # For RNTuples +from podio.arrow_io import Writer as ArrowWriter writer = Writer(filename) # For TTree output (.root) rntuple_writer = RNTupleWriter(filename) # For RNTuple output (.root) +arrow_writer = ArrowWriter(directory) # For Arrow/Parquet output (.podio_parquet) writer.write_frame(frame, category) +arrow_writer.write_frame(frame, category) ``` diff --git a/doc/storage_details.md b/doc/storage_details.md index e0d497df2..8aa765327 100644 --- a/doc/storage_details.md +++ b/doc/storage_details.md @@ -104,3 +104,24 @@ this record. Schematically an SIO file written by podio looks like this SIO file layout schematic + +## Arrow/Parquet + +The Arrow/Parquet backend stores a podio dataset as a directory, usually using +the `.podio_parquet` extension. The directory contains one Parquet file per +category and one `metadata.json` file with dataset-level metadata. + +For a category named `events`, the category data is stored in `events.parquet`. +Each Frame in that category corresponds to one row in the Parquet file. Each +collection is stored as one Arrow column, and Frame parameters are stored in a +special `frame_parameters` column. + +The `metadata.json` file records the podio format marker, the podio version, +the available categories, the category file names, the number of entries, and +the datamodel definitions needed to read the stored collections. + +```{note} +For Arrow/Parquet output all entries of a category have to have the same +collection contents. This content is defined by the first entry that is written +for a category. +``` diff --git a/include/podio/ArrowReader.h b/include/podio/ArrowReader.h new file mode 100644 index 000000000..c7baa47e9 --- /dev/null +++ b/include/podio/ArrowReader.h @@ -0,0 +1,70 @@ +#ifndef PODIO_ARROWREADER_H +#define PODIO_ARROWREADER_H + +#include "podio/podioVersion.h" +#include "podio/utilities/ArrowFrameData.h" +#include "podio/utilities/ReaderCommon.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace arrow { +class Table; +} + +namespace podio { + +/// Arrow backend reader for PODIO +/// +/// Reads data from a directory structure containing one Parquet file per category +/// and a metadata.json file containing metadata for reading. +class ArrowReader : public ReaderCommon { +public: + /// Create an ArrowReader + ArrowReader(); + + /// Open the passed directory for reading. + /// + /// @param directory The path to the directory to read from + void openFile(const std::string& directory); + + ~ArrowReader() = default; + + ArrowReader(const ArrowReader&) = delete; + ArrowReader& operator=(const ArrowReader&) = delete; + ArrowReader(ArrowReader&&) = delete; + ArrowReader& operator=(ArrowReader&&) = delete; + + /// Read the next entry for the given category + std::unique_ptr readNextEntry(std::string_view name, + const std::vector& collsToRead = {}); + + /// Read the specific entry for the given category + std::unique_ptr readEntry(std::string_view name, size_t index, + const std::vector& collsToRead = {}); + + /// Get the number of entries for a category + size_t getEntries(std::string_view name) const; + +private: + struct CategoryInfo { + std::string filePath{}; + size_t entries = 0; + size_t currentIndex = 0; + std::shared_ptr table{nullptr}; + }; + + void loadCategoryTable(CategoryInfo& catInfo); + + std::string m_directory{}; + std::map m_categories{}; +}; + +} // namespace podio + +#endif // PODIO_ARROWREADER_H diff --git a/include/podio/ArrowWriter.h b/include/podio/ArrowWriter.h new file mode 100644 index 000000000..ad577d07a --- /dev/null +++ b/include/podio/ArrowWriter.h @@ -0,0 +1,107 @@ +#ifndef PODIO_ARROWWRITER_H +#define PODIO_ARROWWRITER_H + +#include "podio/utilities/DatamodelRegistryIOHelpers.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +// Forward declarations for Arrow and Parquet +namespace arrow { +class Schema; +class Table; +} // namespace arrow + +namespace podio { + +class Frame; + +/// Arrow backend writer for PODIO +/// +/// Writes data to a directory structure containing one Parquet file per category +/// and a metadata.json file containing metadata for reading. +class ArrowWriter { +public: + /// Configure the ArrowWriter + struct Options { + size_t maxBufferedRows = 1000; + std::string compression = ""; + }; + + /// Create a ArrowWriter to write to a directory. + /// + /// @note Will create the directory if it doesn't exist. Will throw if it exists + /// and is not empty. + /// + /// @param directory The path to the output directory. + /// @param options Configuration options for buffering and compression. + ArrowWriter(const std::string& directory, const Options& options); + explicit ArrowWriter(const std::string& directory); + + /// Destructor writes metadata and closes files. + ~ArrowWriter(); + + ArrowWriter(const ArrowWriter&) = delete; + ArrowWriter& operator=(const ArrowWriter&) = delete; + ArrowWriter(ArrowWriter&&) = delete; + ArrowWriter& operator=(ArrowWriter&&) = delete; + + /// Store the given frame with the given category. + /// + /// @note All frames of the same category must have the same collection + /// names and schemas. Trying to write a frame with different + /// collections to an existing category will result in an exception. + void writeFrame(const podio::Frame& frame, std::string_view category); + + /// Store the given Frame with the given category, specifying collections. + /// + /// @note All frames of the same category must have the same collection + /// names and schemas. Trying to write a frame with different + /// collections to an existing category will result in an exception. + void writeFrame(const podio::Frame& frame, std::string_view category, const std::vector& collsToWrite); + + /// Write the current directory including metadata.json and close files. + void finish(); + +private: + /// Helper struct to manage category state + struct CategoryInfo { + std::string filePath{}; + std::shared_ptr schema{nullptr}; + std::vector collsToWrite{}; + std::vector collTypes{}; + std::vector collIsSubset{}; + std::vector collSchemaVersions{}; + std::vector collIDs{}; + std::vector> buffer{}; + std::unique_ptr writer{nullptr}; + size_t entries = 0; + + CategoryInfo() = default; + ~CategoryInfo() = default; + CategoryInfo(CategoryInfo&&) = default; + CategoryInfo& operator=(CategoryInfo&&) = default; + }; + + void flushCategory(CategoryInfo& catInfo); + void writeMetadata(); + void validateSchema(const CategoryInfo& catInfo, const podio::Frame& frame, + const std::vector& collsToWrite); + + std::string m_directory{}; + Options m_options{}; + std::map m_categories{}; + DatamodelDefinitionCollector m_datamodelCollector{}; + bool m_finished = false; +}; + +} // namespace podio + +#endif // PODIO_ARROWWRITER_H diff --git a/include/podio/utilities/ArrowFrameData.h b/include/podio/utilities/ArrowFrameData.h index df6edfa7b..0626fb5d4 100644 --- a/include/podio/utilities/ArrowFrameData.h +++ b/include/podio/utilities/ArrowFrameData.h @@ -18,7 +18,8 @@ namespace podio { class ArrowFrameData { public: - ArrowFrameData(std::shared_ptr table, int64_t rowIndex); + ArrowFrameData(std::shared_ptr table, int64_t rowIndex, + const std::vector& collsToRead = {}); podio::CollectionIDTable getIDTable() const; std::optional getCollectionBuffers(const std::string& name); diff --git a/include/podio/utilities/ArrowTypeRegistry.h b/include/podio/utilities/ArrowTypeRegistry.h index ffde424a9..da8221849 100644 --- a/include/podio/utilities/ArrowTypeRegistry.h +++ b/include/podio/utilities/ArrowTypeRegistry.h @@ -48,8 +48,7 @@ class ArrowTypeRegistry { std::shared_ptr getType(const std::string& typeName) const; private: - ArrowTypeRegistry() : m_registry() { - } + ArrowTypeRegistry(); std::unordered_map> m_registry; }; diff --git a/podioVersion.in.h b/podioVersion.in.h index 970a0baec..83ac51cb8 100644 --- a/podioVersion.in.h +++ b/podioVersion.in.h @@ -2,8 +2,10 @@ #define PODIO_PODIOVERSION_H #include +#include #include #include +#include #include // Some preprocessor constants and macros for the use cases where they might be @@ -61,6 +63,16 @@ struct Version { return ss.str(); } + static std::optional fromString(const std::string& versionStr) { + uint16_t major = 0, minor = 0, patch = 0; + char dot1, dot2; + std::stringstream ss(versionStr); + if (ss >> major >> dot1 >> minor >> dot2 >> patch && dot1 == '.' && dot2 == '.') { + return Version{major, minor, patch}; + } + return std::nullopt; + } + friend std::ostream& operator<<(std::ostream&, const Version& v); }; diff --git a/python/podio/arrow_io.py b/python/podio/arrow_io.py new file mode 100644 index 000000000..1f7a31e5e --- /dev/null +++ b/python/podio/arrow_io.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +"""Python module for reading arrow files containing podio Frames""" + +from ROOT import gSystem + +if gSystem.DynamicPathName("libpodioArrow.so", True): + gSystem.Load("libpodioArrow") # noqa: 402 +else: + raise ImportError("Error when importing libpodioArrow") +from ROOT import podio # noqa: 402 # pylint: disable=wrong-import-position + +from podio.base_reader import BaseReaderMixin # pylint: disable=wrong-import-position +from podio.base_writer import BaseWriterMixin # pylint: disable=wrong-import-position +from podio.utils import convert_to_str_paths # pylint: disable=wrong-import-position # noqa: E402 + + +class Reader(BaseReaderMixin): + """Reader class for reading podio Arrow files.""" + + def __init__(self, directory): + """Create a reader that reads from the passed directory. + + Args: + directory (str or Path): Directory to open and read data from. + """ + directory = convert_to_str_paths(directory)[0] + self._reader = podio.ArrowReader() + self._reader.openFile(directory) + + super().__init__() + + +class Writer(BaseWriterMixin): + """Writer class for writing podio arrow files.""" + + def __init__(self, directory): + """Create a writer for writing files. + + Args: + directory (str or Path): The name of the output directory. + """ + directory = convert_to_str_paths(directory)[0] + self._writer = podio.ArrowWriter(directory) + + super().__init__() diff --git a/python/templates/ArrowMapper.cc.jinja2 b/python/templates/ArrowMapper.cc.jinja2 index e99d49bca..1359fbdc3 100644 --- a/python/templates/ArrowMapper.cc.jinja2 +++ b/python/templates/ArrowMapper.cc.jinja2 @@ -57,11 +57,11 @@ auto arr_{{ path_prefix }}{{ member.name }} = std::static_pointer_cast<{{ array_ for (size_t arr_i = 0; arr_i < {{ member.array_size }}; ++arr_i) { {% if member.value_builder.is_primitive -%} {% set value = (obj_expr + '.' + member.name + '[arr_i]') -%} - {{ value }} = arr_{{ path_prefix }}{{ member.name }}_val->Value({{ index_var }} * {{ member.array_size }} + arr_i); + {{ value }} = arr_{{ path_prefix }}{{ member.name }}_val->Value((arr_{{ path_prefix }}{{ member.name }}->offset() + {{ index_var }}) * {{ member.array_size }} + arr_i); {% else -%} {% set sub_obj = (obj_expr + '.' + member.name + '[arr_i]') -%} {% for child in member.value_builder.children -%} -{{ read_member(child, sub_obj, path_prefix + member.name + '_val_', index_var) | indent(2, first=True) }} +{{ read_member(child, sub_obj, path_prefix + member.name + '_val_', '(arr_' + path_prefix + member.name + '->offset() + ' ~ index_var ~ ') * ' ~ member.array_size ~ ' + arr_i') | indent(2, first=True) }} {% endfor -%} {% endif -%} } diff --git a/src/ArrowConverterRegistry.cc b/src/ArrowConverterRegistry.cc index 97a208d5e..a7e366079 100644 --- a/src/ArrowConverterRegistry.cc +++ b/src/ArrowConverterRegistry.cc @@ -1,10 +1,77 @@ #include "podio/utilities/ArrowConverterRegistry.h" +#include "podio/CollectionBufferFactory.h" +#include "podio/UserDataCollection.h" +#include "podio/utilities/ArrowTypeRegistry.h" +#include "podio/utilities/ArrowUtils.h" #include "podio/utilities/BackendLibraryLoader.h" +#include namespace podio { +template +void registerPrimitiveConverter(ArrowConverterRegistry& registry) { + const std::string typeName = userDataTypeName(); + registry.registerConverter(typeName, [typeName](const podio::CollectionBase* coll) { + const auto* concreteColl = static_cast*>(coll); + auto type = podio::ArrowTypeRegistry::instance().getType(typeName); + + std::unique_ptr builder; + auto status = arrow::MakeBuilder(arrow::default_memory_pool(), type, &builder); + if (!status.ok()) { + throw std::runtime_error("Failed to create builder for primitive type " + typeName); + } + auto* collectionBuilder = static_cast(builder.get()); + auto* valueBuilder = static_cast(collectionBuilder->value_builder()); + + arrow_utils::checkStatus(collectionBuilder->Append(), "Failed to append to collectionBuilder"); + for (const auto& val : concreteColl->vec()) { + arrow_utils::checkStatus(valueBuilder->Append(val), "Failed to append primitive value"); + } + std::shared_ptr array; + arrow_utils::checkStatus(collectionBuilder->Finish(&array), "Failed to finish collectionBuilder"); + return array; + }); + + registry.registerReader(typeName, + [](const std::shared_ptr& array, int64_t rowIndex, bool isSubset, + podio::SchemaVersionT version) -> std::optional { + auto buffers = podio::CollectionBufferFactory::instance().createBuffers( + std::string(podio::userDataCollTypeName()), version, isSubset); + if (!buffers) { + return std::nullopt; + } + + auto list_array = std::static_pointer_cast(array); + auto obj_array = list_array->value_slice(rowIndex); + auto val_array = std::static_pointer_cast(obj_array); + + auto* dataVec = buffers->dataAsVector(); + size_t collection_size = val_array->length(); + dataVec->reserve(collection_size); + for (size_t i = 0; i < collection_size; ++i) { + dataVec->push_back(val_array->Value(i)); + } + + return buffers; + }); +} + ArrowConverterRegistry& ArrowConverterRegistry::mutInstance() { static ArrowConverterRegistry registry; + static bool registered = false; + if (!registered) { + registered = true; + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + registerPrimitiveConverter(registry); + } return registry; } diff --git a/src/ArrowFrameData.cc b/src/ArrowFrameData.cc index 8848a9c41..63c38ccea 100644 --- a/src/ArrowFrameData.cc +++ b/src/ArrowFrameData.cc @@ -59,7 +59,8 @@ namespace { } // namespace -ArrowFrameData::ArrowFrameData(std::shared_ptr table, int64_t rowIndex) : +ArrowFrameData::ArrowFrameData(std::shared_ptr table, int64_t rowIndex, + const std::vector& collsToRead) : m_table(std::move(table)), m_rowIndex(rowIndex), m_availableCollections(), m_idTable() { if (!m_table) { throw std::runtime_error("ArrowTable is null"); @@ -71,13 +72,26 @@ ArrowFrameData::ArrowFrameData(std::shared_ptr table, int64_t rowI std::vector ids; std::vector names; + if (!collsToRead.empty()) { + auto missing_coll = std::find_if(collsToRead.begin(), collsToRead.end(), [this](const std::string& coll) { + return m_table->GetColumnByName(coll) == nullptr; + }); + if (missing_coll != collsToRead.end()) { + throw std::runtime_error("Collection '" + *missing_coll + "' not found in category."); + } + m_availableCollections = collsToRead; + } + auto schema = m_table->schema(); for (int i = 0; i < schema->num_fields(); ++i) { auto field = schema->field(i); if (field->name() == "frame_parameters") { continue; } - m_availableCollections.push_back(field->name()); + + if (collsToRead.empty()) { + m_availableCollections.push_back(field->name()); + } auto metadata = field->metadata(); if (!metadata) { @@ -100,6 +114,10 @@ ArrowFrameData::ArrowFrameData(std::shared_ptr table, int64_t rowI } std::optional ArrowFrameData::getCollectionBuffers(const std::string& name) { + if (std::find(m_availableCollections.begin(), m_availableCollections.end(), name) == m_availableCollections.end()) { + return std::nullopt; + } + auto chunked_array = m_table->GetColumnByName(name); if (!chunked_array) { return std::nullopt; diff --git a/src/ArrowReader.cc b/src/ArrowReader.cc new file mode 100644 index 000000000..7f3783c5b --- /dev/null +++ b/src/ArrowReader.cc @@ -0,0 +1,172 @@ +#include "podio/ArrowReader.h" + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace podio { + +ArrowReader::ArrowReader() = default; + +void ArrowReader::openFile(const std::string& directory) { + m_directory = directory; + + if (!std::filesystem::exists(m_directory) || !std::filesystem::is_directory(m_directory)) { + throw std::runtime_error("Directory does not exist: " + directory); + } + + auto metadataPath = std::filesystem::path(m_directory) / "metadata.json"; + if (!std::filesystem::exists(metadataPath)) { + throw std::runtime_error("Missing metadata.json in directory: " + directory); + } + + std::ifstream in(metadataPath); + nlohmann::json metadata; + in >> metadata; + + if (metadata.value("format", "") != "podio-arrow") { + throw std::runtime_error("Unsupported format in metadata.json: " + metadata.value("format", "")); + } + if (metadata.value("format_version", 0) != 1) { + throw std::runtime_error("Unsupported format_version in metadata.json: " + + std::to_string(metadata.value("format_version", 0))); + } + + auto versionStr = metadata.value("podio_version", ""); + auto parsedVersion = podio::version::Version::fromString(versionStr); + if (parsedVersion) { + m_fileVersion = parsedVersion.value(); + } else { + throw std::runtime_error("Invalid or missing podio_version in metadata.json: " + versionStr); + } + + for (auto& [name, catJson] : metadata["categories"].items()) { + CategoryInfo catInfo; + catInfo.filePath = (std::filesystem::path(m_directory) / catJson["file"].get()).string(); + catInfo.entries = catJson["entries"].get(); + + m_categories[name] = std::move(catInfo); + m_availableCategories.push_back(name); + } + + std::vector> defs; + std::vector> versions; + + if (metadata.contains("datamodel_definitions")) { + for (auto& [name, def] : metadata["datamodel_definitions"].items()) { + defs.emplace_back(name, def.get()); + } + } + + if (metadata.contains("datamodel_versions")) { + for (auto& [name, versionJson] : metadata["datamodel_versions"].items()) { + versions.emplace_back(name, + podio::version::Version{versionJson["major"].get(), + versionJson["minor"].get(), + versionJson["patch"].get()}); + } + } + + m_datamodelHolder = DatamodelDefinitionHolder(std::move(defs), std::move(versions)); +} + +void ArrowReader::loadCategoryTable(CategoryInfo& catInfo) { + if (catInfo.table) { + return; + } + + if (!std::filesystem::exists(catInfo.filePath)) { + throw std::runtime_error("Missing category file: " + catInfo.filePath); + } + + std::shared_ptr infile; + auto file_result = arrow::io::ReadableFile::Open(catInfo.filePath); + if (!file_result.ok()) { + throw std::runtime_error("Failed to open file: " + file_result.status().ToString()); + } + infile = file_result.ValueOrDie(); + + std::unique_ptr reader; +#if ARROW_VERSION_MAJOR >= 19 + auto reader_result = parquet::arrow::OpenFile(infile, arrow::default_memory_pool()); + if (!reader_result.ok()) { + throw std::runtime_error("Failed to open parquet reader: " + reader_result.status().ToString()); + } + reader = std::move(reader_result.ValueOrDie()); +#else + auto reader_status = parquet::arrow::OpenFile(infile, arrow::default_memory_pool(), &reader); + if (!reader_status.ok()) { + throw std::runtime_error("Failed to open parquet reader: " + reader_status.ToString()); + } +#endif + +#if ARROW_VERSION_MAJOR >= 24 + auto result = reader->ReadTable(); + if (!result.ok()) { + throw std::runtime_error("Failed to read arrow table: " + result.status().ToString()); + } + catInfo.table = std::move(result.ValueOrDie()); +#else + std::shared_ptr table; + auto status = reader->ReadTable(&table); + if (!status.ok()) { + throw std::runtime_error("Failed to read arrow table: " + status.ToString()); + } + catInfo.table = std::move(table); +#endif +} + +std::unique_ptr ArrowReader::readNextEntry(std::string_view name, + const std::vector& collsToRead) { + auto it = m_categories.find(std::string(name)); + if (it == m_categories.end()) { + return nullptr; + } + + if (it->second.currentIndex >= it->second.entries) { + return nullptr; + } + + return readEntry(name, it->second.currentIndex++, collsToRead); +} + +std::unique_ptr ArrowReader::readEntry(std::string_view name, size_t index, + const std::vector& collsToRead) { + auto it = m_categories.find(std::string(name)); + if (it == m_categories.end()) { + return nullptr; + } + + if (index >= it->second.entries) { + return nullptr; + } + + it->second.currentIndex = index + 1; + loadCategoryTable(it->second); + + if (!collsToRead.empty()) { + for (const auto& collName : collsToRead) { + if (it->second.table->schema()->GetFieldIndex(collName) == -1) { + throw std::invalid_argument(collName + " is not available from Frame"); + } + } + } + + return std::make_unique(it->second.table, index, collsToRead); +} + +size_t ArrowReader::getEntries(std::string_view name) const { + auto it = m_categories.find(std::string(name)); + if (it != m_categories.end()) { + return it->second.entries; + } + return 0; +} + +} // namespace podio diff --git a/src/ArrowTypeRegistry.cc b/src/ArrowTypeRegistry.cc index b7becec2e..6e9aee84f 100644 --- a/src/ArrowTypeRegistry.cc +++ b/src/ArrowTypeRegistry.cc @@ -1,8 +1,22 @@ #include "podio/utilities/ArrowTypeRegistry.h" #include "podio/utilities/ArrowConverterRegistry.h" +#include namespace podio { +ArrowTypeRegistry::ArrowTypeRegistry() : m_registry() { + m_registry["int"] = arrow::list(arrow::int32()); + m_registry["float"] = arrow::list(arrow::float32()); + m_registry["double"] = arrow::list(arrow::float64()); + m_registry["uint64_t"] = arrow::list(arrow::uint64()); + m_registry["uint32_t"] = arrow::list(arrow::uint32()); + m_registry["int64_t"] = arrow::list(arrow::int64()); + m_registry["int16_t"] = arrow::list(arrow::int16()); + m_registry["uint16_t"] = arrow::list(arrow::uint16()); + m_registry["int8_t"] = arrow::list(arrow::int8()); + m_registry["uint8_t"] = arrow::list(arrow::uint8()); +} + ArrowTypeRegistry& ArrowTypeRegistry::mutInstance() { static ArrowTypeRegistry registry; return registry; diff --git a/src/ArrowWriter.cc b/src/ArrowWriter.cc new file mode 100644 index 000000000..c3b40e578 --- /dev/null +++ b/src/ArrowWriter.cc @@ -0,0 +1,248 @@ +#include "podio/ArrowWriter.h" + +#include "podio/CollectionBase.h" +#include "podio/DatamodelRegistry.h" +#include "podio/Frame.h" +#include "podio/podioVersion.h" +#include "podio/utilities/ArrowFrameConverter.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#ifndef PODIO_ARROW_DEFAULT_COMPRESSION + #define PODIO_ARROW_DEFAULT_COMPRESSION "UNCOMPRESSED" +#endif + +namespace podio { + +static constexpr std::string_view defaultArrowCompression = PODIO_ARROW_DEFAULT_COMPRESSION; + +ArrowWriter::ArrowWriter(const std::string& directory, const Options& options) : + m_directory(directory), m_options(options) { + if (m_options.compression.empty()) { + m_options.compression = std::string(defaultArrowCompression); + } + if (std::filesystem::exists(m_directory)) { + std::filesystem::remove_all(m_directory); + } + if (m_options.compression != "ZSTD" && m_options.compression != "SNAPPY" && m_options.compression != "UNCOMPRESSED") { + throw std::invalid_argument("Unknown compression: " + m_options.compression); + } + std::filesystem::create_directories(m_directory); +} + +ArrowWriter::ArrowWriter(const std::string& directory) : ArrowWriter(directory, Options{}) { +} + +ArrowWriter::~ArrowWriter() { + try { + finish(); + } catch (const std::exception& e) { + std::cerr << "Exception caught in ArrowWriter destructor: " << e.what() << "\n"; + } +} + +void ArrowWriter::writeFrame(const podio::Frame& frame, std::string_view category) { + writeFrame(frame, category, frame.getAvailableCollections()); +} + +void ArrowWriter::writeFrame(const podio::Frame& frame, std::string_view category, + const std::vector& collsToWrite) { + std::string catName(category); + auto it = m_categories.find(catName); + + std::vector sortedColls = collsToWrite; + std::sort(sortedColls.begin(), sortedColls.end()); + + if (it == m_categories.end()) { + CategoryInfo catInfo; + catInfo.filePath = (std::filesystem::path(m_directory) / (catName + ".parquet")).string(); + catInfo.collsToWrite = sortedColls; + + for (const auto& name : sortedColls) { + const auto* coll = frame.get(name); + if (!coll) { + throw std::runtime_error("Collection " + name + " not found in frame."); + } + catInfo.collTypes.emplace_back(coll->getValueTypeName()); + catInfo.collIsSubset.push_back(coll->isSubsetCollection()); + catInfo.collSchemaVersions.push_back(coll->getSchemaVersion()); + catInfo.collIDs.push_back(coll->getID()); + m_datamodelCollector.registerDatamodelDefinition(coll, name); + } + + m_categories.emplace(catName, std::move(catInfo)); + it = m_categories.find(catName); + } else { + validateSchema(it->second, frame, sortedColls); + } + + auto table = podio::convertFrameToTable(frame, sortedColls); + if (!it->second.schema) { + it->second.schema = table->schema(); + } else if (!it->second.schema->Equals(*table->schema())) { + throw std::runtime_error( + "Arrow schema drift detected: The internal Arrow schema differs for subsequent frames in the same category."); + } + + it->second.buffer.push_back(table); + it->second.entries++; + + if (it->second.buffer.size() >= m_options.maxBufferedRows) { + flushCategory(it->second); + } +} + +void ArrowWriter::validateSchema(const CategoryInfo& catInfo, const podio::Frame& frame, + const std::vector& collsToWrite) { + if (catInfo.collsToWrite != collsToWrite) { + throw std::runtime_error("Schema drift detected: collection names differ for subsequent frames."); + } + + for (size_t i = 0; i < collsToWrite.size(); ++i) { + const auto* coll = frame.get(collsToWrite[i]); + if (!coll) { + throw std::runtime_error("Collection " + collsToWrite[i] + " not found in frame."); + } + if (catInfo.collTypes[i] != coll->getValueTypeName()) { + throw std::runtime_error("Type drift detected for collection " + collsToWrite[i] + ": expected " + + catInfo.collTypes[i] + ", got " + std::string(coll->getValueTypeName())); + } + if (catInfo.collIsSubset[i] != coll->isSubsetCollection()) { + throw std::runtime_error("Subset drift detected for collection " + collsToWrite[i]); + } + if (catInfo.collSchemaVersions[i] != coll->getSchemaVersion()) { + throw std::runtime_error("Schema version drift detected for collection " + collsToWrite[i] + ": expected " + + std::to_string(catInfo.collSchemaVersions[i]) + ", got " + + std::to_string(coll->getSchemaVersion())); + } + if (catInfo.collIDs[i] != coll->getID()) { + throw std::runtime_error("Collection ID drift detected for collection " + collsToWrite[i] + ": expected " + + std::to_string(catInfo.collIDs[i]) + ", got " + std::to_string(coll->getID())); + } + } +} + +void ArrowWriter::flushCategory(CategoryInfo& catInfo) { + if (catInfo.buffer.empty()) { + return; + } + + auto result = arrow::ConcatenateTables(catInfo.buffer); + if (!result.ok()) { + throw std::runtime_error("Failed to concatenate arrow tables: " + result.status().ToString()); + } + const auto& table = result.ValueOrDie(); + + if (!catInfo.writer) { + std::shared_ptr outfile; + auto file_result = arrow::io::FileOutputStream::Open(catInfo.filePath); + if (!file_result.ok()) { + throw std::runtime_error("Failed to open file: " + catInfo.filePath); + } + outfile = file_result.ValueOrDie(); + + parquet::WriterProperties::Builder builder; + if (m_options.compression == "ZSTD") { + builder.compression(parquet::Compression::ZSTD); + } else if (m_options.compression == "SNAPPY") { + builder.compression(parquet::Compression::SNAPPY); + } else { + builder.compression(parquet::Compression::UNCOMPRESSED); + } + + auto arrow_props = parquet::ArrowWriterProperties::Builder().store_schema()->build(); + auto writer_result = parquet::arrow::FileWriter::Open(*catInfo.schema, arrow::default_memory_pool(), outfile, + builder.build(), arrow_props); + if (!writer_result.ok()) { + throw std::runtime_error("Failed to open parquet writer: " + writer_result.status().ToString()); + } + catInfo.writer = std::move(writer_result.ValueOrDie()); + } + + auto status = catInfo.writer->WriteTable(*table, table->num_rows()); + if (!status.ok()) { + throw std::runtime_error("Failed to write table to parquet: " + status.ToString()); + } + catInfo.buffer.clear(); +} + +void ArrowWriter::finish() { + if (m_finished) { + return; + } + + for (auto& [name, catInfo] : m_categories) { + flushCategory(catInfo); + if (catInfo.writer) { + auto status = catInfo.writer->Close(); + if (!status.ok()) { + throw std::runtime_error("Failed to close parquet writer: " + status.ToString()); + } + } + } + + writeMetadata(); + m_finished = true; +} + +void ArrowWriter::writeMetadata() { + nlohmann::json metadata; + metadata["format"] = "podio-arrow"; + metadata["format_version"] = 1; + metadata["podio_version"] = std::string(podio::version::build_version); + + nlohmann::json categoriesJson; + for (const auto& [name, catInfo] : m_categories) { + nlohmann::json catJson; + catJson["file"] = std::filesystem::path(catInfo.filePath).filename().string(); + catJson["entries"] = catInfo.entries; + + nlohmann::json collectionsJson = nlohmann::json::array(); + for (size_t i = 0; i < catInfo.collsToWrite.size(); ++i) { + nlohmann::json collJson; + collJson["name"] = catInfo.collsToWrite[i]; + collJson["value_type"] = catInfo.collTypes[i]; + collJson["schema_version"] = catInfo.collSchemaVersions[i]; + collJson["is_subset"] = catInfo.collIsSubset[i]; + collJson["id"] = catInfo.collIDs[i]; + collectionsJson.push_back(collJson); + } + catJson["collections"] = collectionsJson; + categoriesJson[name] = catJson; + } + metadata["categories"] = categoriesJson; + nlohmann::json edmDefsJson = nlohmann::json::object(); + nlohmann::json edmVersionsJson = nlohmann::json::object(); + for (const auto& [name, def] : m_datamodelCollector.getDatamodelDefinitionsToWrite()) { + edmDefsJson[name] = def; + auto edmVersion = podio::DatamodelRegistry::instance().getDatamodelVersion(name); + if (edmVersion) { + edmVersionsJson[name] = {{"major", edmVersion.value().major}, + {"minor", edmVersion.value().minor}, + {"patch", edmVersion.value().patch}}; + } + } + metadata["datamodel_definitions"] = edmDefsJson; + metadata["datamodel_versions"] = edmVersionsJson; + + auto finalPath = std::filesystem::path(m_directory) / "metadata.json"; + + std::ofstream out(finalPath); + if (!out) { + throw std::runtime_error("Failed to open metadata.json for writing"); + } + out << metadata.dump(2); + out.close(); +} + +} // namespace podio diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 254803544..3a52679c4 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -160,16 +160,48 @@ if(ENABLE_ARROW) ArrowFrameData.cc ) - add_library(podioArrow SHARED ${arrow_sources}) - add_library(podio::podioArrow ALIAS podioArrow) - target_include_directories(podioArrow PUBLIC - $ - $ - ) - target_link_libraries(podioArrow PUBLIC podio::podio ${PODIO_ARROW_TARGET}) + set(arrow_headers) + if(ENABLE_PARQUET) + LIST(APPEND arrow_sources + ArrowWriter.cc + ArrowReader.cc + ) + LIST(APPEND arrow_headers + "${PROJECT_SOURCE_DIR}/include/podio/ArrowReader.h" + "${PROJECT_SOURCE_DIR}/include/podio/ArrowWriter.h" + ) + set(arrow_selection "arrow_selection.xml") + PODIO_ADD_LIB_AND_DICT(podioArrow "${arrow_headers}" "${arrow_sources}" ${arrow_selection}) + target_link_libraries(podioArrow PUBLIC podio::podio ${PODIO_ARROW_TARGET}) + target_link_libraries(podioArrow PRIVATE nlohmann_json::nlohmann_json) + target_link_libraries(podioArrow PUBLIC ${PODIO_PARQUET_TARGET}) + target_compile_definitions(podioArrow PUBLIC PODIO_ENABLE_PARQUET=1) + else() + add_library(podioArrow SHARED ${arrow_sources}) + add_library(podio::podioArrow ALIAS podioArrow) + target_include_directories(podioArrow PUBLIC + $ + $ + ) + target_link_libraries(podioArrow PUBLIC podio::podio ${PODIO_ARROW_TARGET}) + endif() + target_compile_definitions(podioArrow PUBLIC PODIO_ENABLE_ARROW=1) + if (ARROW_WITH_ZSTD) + set(DEFAULT_COMPRESSION_STR "ZSTD") + elseif(ARROW_WITH_SNAPPY) + set(DEFAULT_COMPRESSION_STR "SNAPPY") + else() + set(DEFAULT_COMPRESSION_STR "UNCOMPRESSED") + endif() + set(PODIO_ARROW_DEFAULT_COMPRESSION "${DEFAULT_COMPRESSION_STR}" CACHE STRING "Default compression for Arrow/Parquet") + target_compile_definitions(podioArrow PRIVATE PODIO_ARROW_DEFAULT_COMPRESSION="${PODIO_ARROW_DEFAULT_COMPRESSION}") + LIST(APPEND INSTALL_LIBRARIES podioArrow) + if(ENABLE_PARQUET) + LIST(APPEND INSTALL_LIBRARIES podioArrowDict) + endif() endif() @@ -193,6 +225,9 @@ target_link_libraries(podioIO PUBLIC podio::podio podio::podioRootIO) if(ENABLE_SIO) target_link_libraries(podioIO PUBLIC podio::podioSioIO) endif() +if(ENABLE_ARROW) + target_link_libraries(podioIO PUBLIC podio::podioArrow) +endif() # --- DataSource @@ -277,3 +312,11 @@ if (ENABLE_DATASOURCE) DESTINATION "${CMAKE_INSTALL_LIBDIR}" ) endif() + +if (ENABLE_PARQUET) + install(FILES + ${CMAKE_CURRENT_BINARY_DIR}/podioArrowDictDict.rootmap + ${CMAKE_CURRENT_BINARY_DIR}/libpodioArrowDict_rdict.pcm + DESTINATION "${CMAKE_INSTALL_LIBDIR}" + ) +endif() diff --git a/src/Reader.cc b/src/Reader.cc index 922d62eea..f01614381 100644 --- a/src/Reader.cc +++ b/src/Reader.cc @@ -7,12 +7,16 @@ #if PODIO_ENABLE_SIO #include "podio/SIOReader.h" #endif +#if PODIO_ENABLE_ARROW && PODIO_ENABLE_PARQUET + #include "podio/ArrowReader.h" +#endif #include "podio/utilities/Glob.h" #include "podio/utilities/ReaderUtils.h" #include "TFile.h" #include "TKey.h" +#include #include namespace podio { @@ -83,6 +87,20 @@ Reader makeReader(const std::vector& filenames) { return reader; #else throw std::runtime_error("SIO reader not available. Please recompile with SIO support."); +#endif + } else if (suffix == "podio_parquet" || + (std::filesystem::is_directory(filenames[0]) && + std::filesystem::exists(std::filesystem::path(filenames[0]) / "metadata.json"))) { +#if PODIO_ENABLE_ARROW && PODIO_ENABLE_PARQUET + if (filenames.size() > 1) { + throw std::runtime_error("The Arrow reader does currently not support reading multiple directories"); + } + auto actualReader = std::make_unique(); + actualReader->openFile(filenames[0]); + Reader reader{std::move(actualReader)}; + return reader; +#else + throw std::runtime_error("Arrow reader not available. Please recompile with Arrow and Parquet support."); #endif } diff --git a/src/Writer.cc b/src/Writer.cc index 84bd82a60..128c7dd01 100644 --- a/src/Writer.cc +++ b/src/Writer.cc @@ -7,6 +7,9 @@ #if PODIO_ENABLE_SIO #include "podio/SIOWriter.h" #endif +#if PODIO_ENABLE_ARROW && PODIO_ENABLE_PARQUET + #include "podio/ArrowWriter.h" +#endif #include #include @@ -45,6 +48,12 @@ Writer makeWriter(const std::string& filename, const std::string& type) { return Writer{std::make_unique(filename)}; #else throw std::runtime_error("SIO writer not available. Please recompile with SIO support."); +#endif + } else if (endsWith(filename, ".podio_parquet") || lower(type) == "parquet") { +#if PODIO_ENABLE_ARROW && PODIO_ENABLE_PARQUET + return Writer{std::make_unique(filename)}; +#else + throw std::runtime_error("Arrow writer not available. Please recompile with Arrow and Parquet support."); #endif } throw std::runtime_error("Unknown file type for file " + filename + " with type " + type); diff --git a/src/arrow_selection.xml b/src/arrow_selection.xml new file mode 100644 index 000000000..093457ef9 --- /dev/null +++ b/src/arrow_selection.xml @@ -0,0 +1,4 @@ + + + + diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c0b3f3b85..1d83ab00a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -94,6 +94,11 @@ if (ENABLE_SIO) add_subdirectory(sio_io) endif() + +if (ENABLE_ARROW AND ENABLE_PARQUET) + add_subdirectory(arrow_io) +endif() + add_subdirectory(unittests) add_subdirectory(dumpmodel) add_subdirectory(schema_evolution) diff --git a/tests/arrow_io/CMakeLists.txt b/tests/arrow_io/CMakeLists.txt new file mode 100644 index 000000000..0cd56b2f2 --- /dev/null +++ b/tests/arrow_io/CMakeLists.txt @@ -0,0 +1,38 @@ +set(arrow_dependent_tests + read_frame_arrow.cpp + write_frame_arrow.cpp + read_and_write_frame_arrow.cpp + read_python_frame_arrow.cpp + write_interface_arrow.cpp + read_interface_arrow.cpp + selected_colls_roundtrip_arrow.cpp + write_frame_arrow_multithreaded.cpp + read_frame_arrow_multithreaded.cpp +) +set(arrow_libs podio::podioArrow podio::podioIO TestDataModelPodioArrow ExtensionDataModelPodioArrow InterfaceExtensionDataModelPodioArrow) +foreach( sourcefile ${arrow_dependent_tests} ) + CREATE_PODIO_TEST(${sourcefile} "${arrow_libs}") +endforeach() + +set_tests_properties(write_frame_arrow PROPERTIES FIXTURES_SETUP podio_write_arrow_fixture) +set_tests_properties(write_frame_arrow_multithreaded PROPERTIES FIXTURES_SETUP podio_write_arrow_mt_fixture) +set_tests_properties(write_interface_arrow PROPERTIES FIXTURES_SETUP podio_write_interface_arrow_fixture) + +set_tests_properties( + read_frame_arrow + read_and_write_frame_arrow + selected_colls_roundtrip_arrow + + PROPERTIES + FIXTURES_REQUIRED podio_write_arrow_fixture +) + +set_tests_properties(read_interface_arrow PROPERTIES FIXTURES_REQUIRED podio_write_interface_arrow_fixture) +set_tests_properties(read_frame_arrow_multithreaded PROPERTIES FIXTURES_REQUIRED podio_write_arrow_mt_fixture) + +#--- Write via python and the Arrow backend and see if we can read it back in in +#--- c++ +add_test(NAME write_python_frame_arrow COMMAND python3 ${PROJECT_SOURCE_DIR}/tests/write_frame.py example_frame_with_py.podio_parquet arrow_io.Writer) +PODIO_SET_TEST_ENV(write_python_frame_arrow PYTHON) +set_tests_properties(write_python_frame_arrow PROPERTIES FIXTURES_SETUP podio_write_python_arrow_fixture) +set_tests_properties(read_python_frame_arrow PROPERTIES FIXTURES_REQUIRED podio_write_python_arrow_fixture) diff --git a/tests/arrow_io/read_and_write_frame_arrow.cpp b/tests/arrow_io/read_and_write_frame_arrow.cpp new file mode 100644 index 000000000..f01c59c09 --- /dev/null +++ b/tests/arrow_io/read_and_write_frame_arrow.cpp @@ -0,0 +1,12 @@ +#include "read_and_write_frame.h" + +#include "podio/ArrowReader.h" +#include "podio/ArrowWriter.h" +#include + +int main() { + std::filesystem::remove_all("rewritten_frame.podio_parquet"); + return rewrite_frames("example_frame.podio_parquet", + "rewritten_frame.podio_parquet") + + read_rewritten_frames("rewritten_frame.podio_parquet"); +} diff --git a/tests/arrow_io/read_frame_arrow.cpp b/tests/arrow_io/read_frame_arrow.cpp new file mode 100644 index 000000000..d994f5781 --- /dev/null +++ b/tests/arrow_io/read_frame_arrow.cpp @@ -0,0 +1,16 @@ +#include "read_frame.h" +#include "read_frame_auxiliary.h" + +#include "podio/ArrowReader.h" + +int main(int argc, char* argv[]) { + std::string inputFile = "example_frame.podio_parquet"; + bool assertBuildVersion = true; + if (argc == 2) { + inputFile = argv[1]; + assertBuildVersion = false; + } + + return read_frames(inputFile, assertBuildVersion) + + test_frame_aux_info(inputFile) + test_read_frame_limited(inputFile); +} diff --git a/tests/arrow_io/read_frame_arrow_multithreaded.cpp b/tests/arrow_io/read_frame_arrow_multithreaded.cpp new file mode 100644 index 000000000..a1f06976a --- /dev/null +++ b/tests/arrow_io/read_frame_arrow_multithreaded.cpp @@ -0,0 +1,20 @@ +#include "read_frame_multithreaded.h" + +#include "podio/ArrowReader.h" + +#include + +int main(int argc, char* argv[]) { + int nThreads = 4; + int framesPerThread = 10; + if (argc >= 2) { + nThreads = std::atoi(argv[1]); + } + if (argc >= 3) { + framesPerThread = std::atoi(argv[2]); + } + + const unsigned expectedEntries = nThreads * framesPerThread; + return read_frames_multithreaded("example_frame_arrow_multithreaded.podio_parquet", nThreads, + expectedEntries); +} diff --git a/tests/arrow_io/read_interface_arrow.cpp b/tests/arrow_io/read_interface_arrow.cpp new file mode 100644 index 000000000..67fc4f44b --- /dev/null +++ b/tests/arrow_io/read_interface_arrow.cpp @@ -0,0 +1,7 @@ +#include "read_frame.h" +#include "read_interface.h" + +int main(int, char**) { + auto readerArrow = podio::makeReader("example_frame_arrow_interface.podio_parquet"); + return read_frames(readerArrow) + test_read_frame_limited(readerArrow); +} diff --git a/tests/arrow_io/read_python_frame_arrow.cpp b/tests/arrow_io/read_python_frame_arrow.cpp new file mode 100644 index 000000000..e8ecb0e57 --- /dev/null +++ b/tests/arrow_io/read_python_frame_arrow.cpp @@ -0,0 +1,7 @@ +#include "read_python_frame.h" + +#include "podio/ArrowReader.h" + +int main() { + return read_frame("example_frame_with_py.podio_parquet"); +} diff --git a/tests/arrow_io/selected_colls_roundtrip_arrow.cpp b/tests/arrow_io/selected_colls_roundtrip_arrow.cpp new file mode 100644 index 000000000..6b67e48d8 --- /dev/null +++ b/tests/arrow_io/selected_colls_roundtrip_arrow.cpp @@ -0,0 +1,11 @@ +#include "selected_colls_roundtrip.h" + +#include "podio/ArrowReader.h" +#include "podio/ArrowWriter.h" +#include + +int main() { + std::filesystem::remove_all("selected_example_frame.podio_parquet"); + return do_roundtrip("example_frame.podio_parquet", + "selected_example_frame.podio_parquet"); +} diff --git a/tests/arrow_io/write_frame_arrow.cpp b/tests/arrow_io/write_frame_arrow.cpp new file mode 100644 index 000000000..4b09ed1d7 --- /dev/null +++ b/tests/arrow_io/write_frame_arrow.cpp @@ -0,0 +1,12 @@ +#include "write_frame.h" + +#include "podio/ArrowWriter.h" + +#include + +int main(int, char**) { + std::string filename = "example_frame.podio_parquet"; + std::filesystem::remove_all(filename); + write_frames(filename); + return 0; +} diff --git a/tests/arrow_io/write_frame_arrow_multithreaded.cpp b/tests/arrow_io/write_frame_arrow_multithreaded.cpp new file mode 100644 index 000000000..7f06294b6 --- /dev/null +++ b/tests/arrow_io/write_frame_arrow_multithreaded.cpp @@ -0,0 +1,21 @@ +#include "write_frame_multithreaded.h" + +#include "podio/ArrowWriter.h" + +#include +#include + +int main(int argc, char* argv[]) { + int nThreads = 4; + int framesPerThread = 10; + if (argc >= 2) { + nThreads = std::atoi(argv[1]); + } + if (argc >= 3) { + framesPerThread = std::atoi(argv[2]); + } + + std::string filename = "example_frame_arrow_multithreaded.podio_parquet"; + std::filesystem::remove_all(filename); + return write_frames_multithreaded(filename, nThreads, framesPerThread); +} diff --git a/tests/arrow_io/write_interface_arrow.cpp b/tests/arrow_io/write_interface_arrow.cpp new file mode 100644 index 000000000..bfb551678 --- /dev/null +++ b/tests/arrow_io/write_interface_arrow.cpp @@ -0,0 +1,12 @@ +#include "write_interface.h" + +#include + +int main(int, char**) { + + std::filesystem::remove_all("example_frame_arrow_interface.podio_parquet"); + auto writerArrow = podio::makeWriter("example_frame_arrow_interface.podio_parquet", "parquet"); + write_frames(writerArrow); + + return 0; +} diff --git a/tests/schema_evolution/code_gen/test_utilities.cmake b/tests/schema_evolution/code_gen/test_utilities.cmake index 726358e0f..215283baf 100644 --- a/tests/schema_evolution/code_gen/test_utilities.cmake +++ b/tests/schema_evolution/code_gen/test_utilities.cmake @@ -30,9 +30,12 @@ function(GENERATE_DATAMODEL test_case model_version) endif() # Generate the datamodel with appropriate options + set(TEST_IO_HANDLERS ${PODIO_IO_HANDLERS}) + list(REMOVE_ITEM TEST_IO_HANDLERS "ARROW") + if(PARSED_ARGS_WITH_EVOLUTION) PODIO_GENERATE_DATAMODEL(datamodel ${test_case}/${model_version}.yaml headers sources - IO_BACKEND_HANDLERS ${PODIO_IO_HANDLERS} + IO_BACKEND_HANDLERS ${TEST_IO_HANDLERS} OUTPUT_FOLDER ${output_base} OLD_DESCRIPTIONS ${old_descriptions} SCHEMA_EVOLUTION ${test_case}/evolution.yaml @@ -40,13 +43,13 @@ function(GENERATE_DATAMODEL test_case model_version) else() if(old_descriptions AND NOT PARSED_ARGS_NO_EVOLUTION_CHECKS) PODIO_GENERATE_DATAMODEL(datamodel ${test_case}/${model_version}.yaml headers sources - IO_BACKEND_HANDLERS ${PODIO_IO_HANDLERS} + IO_BACKEND_HANDLERS ${TEST_IO_HANDLERS} OUTPUT_FOLDER ${output_base} OLD_DESCRIPTIONS ${old_descriptions} ) else() PODIO_GENERATE_DATAMODEL(datamodel ${test_case}/${model_version}.yaml headers sources - IO_BACKEND_HANDLERS ${PODIO_IO_HANDLERS} + IO_BACKEND_HANDLERS ${TEST_IO_HANDLERS} OUTPUT_FOLDER ${output_base} ) endif() diff --git a/tests/write_frame.py b/tests/write_frame.py index c814fbd42..95b02f9a2 100644 --- a/tests/write_frame.py +++ b/tests/write_frame.py @@ -108,5 +108,10 @@ def write_file(writer_type, filename): args = parser.parse_args() io_format = args.outputfile.split(".")[-1] + if io_format == "podio_parquet": + ROOT.gSystem.Load("libpodioArrow") + ROOT.gSystem.Load("libTestDataModelArrow") + ROOT.gSystem.Load("libExtensionDataModelArrow") + ROOT.gSystem.Load("libInterfaceExtensionDataModelArrow") write_file(args.writer, args.outputfile)