diff --git a/.cppcheck-suppressions.txt b/.cppcheck-suppressions.txt index 2efdf1392..bf7345478 100644 --- a/.cppcheck-suppressions.txt +++ b/.cppcheck-suppressions.txt @@ -15,3 +15,6 @@ useStlAlgorithm:LinkCollectionImpl.h useStlAlgorithm:podio-dump-tool.cpp useStlAlgorithm:SIOFrameData.cc useStlAlgorithm:rootUtils.h + +# getAvailableCollections must return by value to satisfy C++20 FrameDataType concept (std::same_as) +returnByReference:*ArrowFrameData* diff --git a/include/podio/utilities/ArrowConverterRegistry.h b/include/podio/utilities/ArrowConverterRegistry.h index 5168687e9..6d3200314 100644 --- a/include/podio/utilities/ArrowConverterRegistry.h +++ b/include/podio/utilities/ArrowConverterRegistry.h @@ -1,10 +1,14 @@ #ifndef PODIO_ARROWCONVERTERREGISTRY_H #define PODIO_ARROWCONVERTERREGISTRY_H +#include "podio/CollectionBuffers.h" +#include "podio/SchemaEvolution.h" +#include #include +#include #include +#include #include -#include namespace arrow { class Array; @@ -20,6 +24,8 @@ class CollectionBase; class ArrowConverterRegistry { public: using CreatorFunc = std::function(const podio::CollectionBase*)>; + using BufferReaderFunc = std::function(std::shared_ptr, + int64_t, bool, SchemaVersionT)>; ArrowConverterRegistry(const ArrowConverterRegistry&) = delete; ArrowConverterRegistry& operator=(const ArrowConverterRegistry&) = delete; @@ -44,11 +50,23 @@ class ArrowConverterRegistry { */ CreatorFunc getConverter(const std::string& typeName) const; + /** + * @brief Register an Arrow array reader callback for a specific type name. + */ + void registerReader(const std::string& typeName, BufferReaderFunc&& reader); + + /** + * @brief Retrieve the Arrow collection reader registered for a specific type name. + * @return The registered BufferReaderFunc, or nullptr if not registered. + */ + BufferReaderFunc getReader(const std::string& typeName) const; + private: - ArrowConverterRegistry() : m_registry() { + ArrowConverterRegistry() : m_registry(), m_readerRegistry() { } - std::unordered_map m_registry; + std::map m_registry; + std::map m_readerRegistry; }; } // namespace podio diff --git a/include/podio/utilities/ArrowFrameConverter.h b/include/podio/utilities/ArrowFrameConverter.h index 1df53d25a..f0f90b2ce 100644 --- a/include/podio/utilities/ArrowFrameConverter.h +++ b/include/podio/utilities/ArrowFrameConverter.h @@ -35,6 +35,17 @@ std::shared_ptr objectRefType(); std::shared_ptr convertFrameToTable(const podio::Frame& frame, const std::vector& collsToWrite); +/** + * @brief Convert a 1-row slice of an Arrow Table to a PODIO Frame. + * + * Reconstructs the frame parameters and all collection buffers using registered BufferReaderFuncs. + * + * @param table The Arrow Table containing the frame data. + * @param rowIndex The index of the row to read. + * @return The reconstructed Frame. + */ +podio::Frame convertTableToFrame(const std::shared_ptr& table, int rowIndex = 0); + } // namespace podio #endif // PODIO_ARROWFRAMECONVERTER_H diff --git a/include/podio/utilities/ArrowFrameData.h b/include/podio/utilities/ArrowFrameData.h new file mode 100644 index 000000000..df6edfa7b --- /dev/null +++ b/include/podio/utilities/ArrowFrameData.h @@ -0,0 +1,38 @@ +#ifndef PODIO_ARROWFRAMEDATA_H +#define PODIO_ARROWFRAMEDATA_H + +#include "podio/CollectionBuffers.h" +#include "podio/CollectionIDTable.h" +#include "podio/GenericParameters.h" + +#include +#include +#include +#include + +namespace arrow { +class Table; +} // namespace arrow + +namespace podio { + +class ArrowFrameData { +public: + ArrowFrameData(std::shared_ptr table, int64_t rowIndex); + + podio::CollectionIDTable getIDTable() const; + std::optional getCollectionBuffers(const std::string& name); + // Must return by value to satisfy the C++20 FrameDataType concept (as other backends construct it dynamically) + std::vector getAvailableCollections() const; + std::unique_ptr getParameters(); + +private: + std::shared_ptr m_table; + int64_t m_rowIndex; + std::vector m_availableCollections; + podio::CollectionIDTable m_idTable; +}; + +} // namespace podio + +#endif // PODIO_ARROWFRAMEDATA_H diff --git a/python/podio_gen/cpp_generator.py b/python/podio_gen/cpp_generator.py index be09754fc..01c2b6545 100644 --- a/python/podio_gen/cpp_generator.py +++ b/python/podio_gen/cpp_generator.py @@ -81,6 +81,34 @@ "std::string": "arrow::StringBuilder", } +ARROW_ARRAYS = { + "bool": "arrow::BooleanArray", + "char": "arrow::Int8Array", + "short": "arrow::Int16Array", + "int": "arrow::Int32Array", + "long": "arrow::Int64Array", + "long long": "arrow::Int64Array", + "unsigned": "arrow::UInt32Array", + "unsigned int": "arrow::UInt32Array", + "unsigned long": "arrow::UInt64Array", + "unsigned long long": "arrow::UInt64Array", + "float": "arrow::FloatArray", + "double": "arrow::DoubleArray", + "int16_t": "arrow::Int16Array", + "int32_t": "arrow::Int32Array", + "int64_t": "arrow::Int64Array", + "uint16_t": "arrow::UInt16Array", + "uint32_t": "arrow::UInt32Array", + "uint64_t": "arrow::UInt64Array", + "std::int16_t": "arrow::Int16Array", + "std::int32_t": "arrow::Int32Array", + "std::int64_t": "arrow::Int64Array", + "std::uint16_t": "arrow::UInt16Array", + "std::uint32_t": "arrow::UInt32Array", + "std::uint64_t": "arrow::UInt64Array", + "std::string": "arrow::StringArray", +} + class IncludeFrom(IntEnum): """Enum to signify if an include is needed and from where it should come""" @@ -194,6 +222,7 @@ def _write_arrow_mapper_header(self, datamodel): "package_name": self.package_name, "schema_version": self.datamodel.schema_version, "datatypes": datatypes, + "links": datamodel.get("links", []), "incfolder": self.incfolder, } self._write_file( @@ -243,6 +272,7 @@ def _get_member_metadata(self, member): "name": member.name, "getter": member.getter_name(self.get_syntax), "is_array": member.is_array, + "full_type": member.full_type, } if member.is_array: meta["array_size"] = member.array_size @@ -250,7 +280,9 @@ def _get_member_metadata(self, member): if member.array_type in ARROW_PRIMITIVE_TYPES: meta["value_builder"] = { "builder_type": ARROW_BUILDERS[member.array_type], + "arrow_array_type": ARROW_ARRAYS[member.array_type], "is_primitive": True, + "full_type": member.array_type, } else: comp_name = member.array_type.removeprefix("::") @@ -259,10 +291,12 @@ def _get_member_metadata(self, member): ) meta["value_builder"] = { "builder_type": "arrow::StructBuilder", + "arrow_array_type": "arrow::StructArray", "is_struct": True, "children": [ self._get_member_metadata(sub_mem) for sub_mem in comp["Members"] ], + "full_type": member.array_type, } return meta @@ -272,10 +306,12 @@ def _get_member_metadata(self, member): ) if comp: meta["builder_type"] = "arrow::StructBuilder" + meta["arrow_array_type"] = "arrow::StructArray" meta["is_struct"] = True meta["children"] = [self._get_member_metadata(sub_mem) for sub_mem in comp["Members"]] else: meta["builder_type"] = ARROW_BUILDERS.get(member.full_type) + meta["arrow_array_type"] = ARROW_ARRAYS.get(member.full_type) meta["is_primitive"] = True return meta @@ -284,11 +320,14 @@ def _get_vector_metadata(self, member): meta = { "name": member.name, "getter": member.getter_name(self.get_syntax), + "full_type": member.full_type, } if member.full_type in ARROW_PRIMITIVE_TYPES: meta["value_builder"] = { "builder_type": ARROW_BUILDERS[member.full_type], + "arrow_array_type": ARROW_ARRAYS[member.full_type], "is_primitive": True, + "full_type": member.full_type, } else: comp_name = member.full_type.removeprefix("::") @@ -297,8 +336,10 @@ def _get_vector_metadata(self, member): ) meta["value_builder"] = { "builder_type": "arrow::StructBuilder", + "arrow_array_type": "arrow::StructArray", "is_struct": True, "children": [self._get_member_metadata(sub_mem) for sub_mem in comp["Members"]], + "full_type": member.full_type, } return meta diff --git a/python/templates/ArrowMapper.cc.jinja2 b/python/templates/ArrowMapper.cc.jinja2 index abe6dbe95..cb6329ca8 100644 --- a/python/templates/ArrowMapper.cc.jinja2 +++ b/python/templates/ArrowMapper.cc.jinja2 @@ -4,6 +4,7 @@ #include "podio/utilities/ArrowFrameConverter.h" #include "podio/utilities/ArrowTypeRegistry.h" #include "podio/utilities/ArrowConverterRegistry.h" +#include "podio/CollectionBufferFactory.h" #include #include @@ -14,15 +15,71 @@ #include #include +#include "podio/LinkCollection.h" + {% for datatype in datatypes %} #include "{{ incfolder }}{{ datatype.class.bare_type }}Collection.h" {% endfor %} +{% for link in links %} +#include "{{ incfolder }}{{ link['class'].bare_type }}Collection.h" +{% endfor %} namespace {{ package_name }}::arrow_io { namespace { using podio::arrow_utils::checkStatus; +{% macro declare_arrays(member, parent_array, path_prefix) -%} +{% if member.is_array -%} +auto arr_{{ path_prefix }}{{ member.name }} = std::static_pointer_cast({{ parent_array }}->GetFieldByName("{{ member.name }}")); +{% set val_array_type = member.value_builder.arrow_array_type -%} +{% if member.value_builder.is_primitive -%} +auto arr_{{ path_prefix }}{{ member.name }}_val = std::static_pointer_cast<{{ val_array_type }}>(arr_{{ path_prefix }}{{ member.name }}->values()); +{% else -%} +auto arr_{{ path_prefix }}{{ member.name }}_val = std::static_pointer_cast(arr_{{ path_prefix }}{{ member.name }}->values()); +{% for child in member.value_builder.children -%} +{{ declare_arrays(child, 'arr_' + path_prefix + member.name + '_val', path_prefix + member.name + '_val_') }} +{% endfor -%} +{% endif -%} +{% elif member.is_struct -%} +auto arr_{{ path_prefix }}{{ member.name }} = std::static_pointer_cast({{ parent_array }}->GetFieldByName("{{ member.name }}")); +{% for child in member.children -%} +{{ declare_arrays(child, 'arr_' + path_prefix + member.name, path_prefix + member.name + '_') }} +{% endfor -%} +{% else -%} +{% set array_type = member.arrow_array_type -%} +auto arr_{{ path_prefix }}{{ member.name }} = std::static_pointer_cast<{{ array_type }}>({{ parent_array }}->GetFieldByName("{{ member.name }}")); +{% endif -%} +{%- endmacro %} + +{% macro read_member(member, obj_expr, path_prefix, index_var='i') -%} +{% if member.is_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); +{% 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) }} +{% endfor -%} +{% endif -%} +} +{% elif member.is_struct -%} +{% set sub_obj = (obj_expr + '.' + member.name) -%} +{% for child in member.children -%} +{{ read_member(child, sub_obj, path_prefix + member.name + '_', index_var) }} +{% endfor -%} +{% else -%} +{% set value = (obj_expr + '.' + member.name) -%} +{% if member.full_type == 'std::string' -%} + {{ value }} = arr_{{ path_prefix }}{{ member.name }}->GetString({{ index_var }}); +{% else -%} + {{ value }} = arr_{{ path_prefix }}{{ member.name }}->Value({{ index_var }}); +{% endif -%} +{% endif -%} +{%- endmacro %} + static const std::map>& typeMap() { static const std::map> instance = { {% for datatype in datatypes %} @@ -34,6 +91,16 @@ static const std::map>& typeMap() {% endfor %} })) }, +{% endfor %} +{% for link in links %} + { + "podio::Link<{{ link.From.full_type }},{{ link.To.full_type }}>", + arrow::list(arrow::struct_({ + arrow::field("weight", arrow::float32()), + arrow::field("from", podio::objectRefType()), + arrow::field("to", podio::objectRefType()), + })) + }, {% endfor %} }; return instance; @@ -62,7 +129,6 @@ bool registerConverters() { } auto* collectionBuilder = static_cast(builder.get()); auto* objectBuilder = static_cast(collectionBuilder->value_builder()); - {% for member in datatype.arrow_metadata.members %} {{ arrow_macros.declare_builder(member, 'objectBuilder', loop.index0, '') | indent(4) }} {%- endfor %} @@ -95,12 +161,206 @@ bool registerConverters() { checkStatus(collectionBuilder->Finish(&array), "Failed to finish collectionBuilder"); return array; }); +{% endfor %} +{% for link in links %} + convRegistry.registerConverter("podio::Link<{{ link.From.full_type }},{{ link.To.full_type }}>", [](const podio::CollectionBase* coll) { + const auto* concreteColl = static_cast*>(coll); + auto type = podio::ArrowTypeRegistry::instance().getType("podio::Link<{{ link.From.full_type }},{{ link.To.full_type }}>"); + 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 Link: " + status.ToString()); + } + auto* collectionBuilder = static_cast(builder.get()); + auto* objectBuilder = static_cast(collectionBuilder->value_builder()); + auto* weightBuilder = static_cast(objectBuilder->child(0)); + auto* fromBuilder = static_cast(objectBuilder->child(1)); + auto* toBuilder = static_cast(objectBuilder->child(2)); + + auto appendRelationRef = [](arrow::StructBuilder* structBuilder, const podio::ObjectID& objId) { + checkStatus(structBuilder->Append(), "Failed to append relation struct"); + checkStatus(static_cast(structBuilder->child(0))->Append(objId.collectionID), "Failed to append collectionID"); + checkStatus(static_cast(structBuilder->child(1))->Append(objId.index), "Failed to append index"); + }; + + checkStatus(collectionBuilder->Append(), "Failed to append to collectionBuilder"); + for (size_t i = 0; i < concreteColl->size(); ++i) { + auto obj = (*concreteColl)[i]; + checkStatus(objectBuilder->Append(), "Failed to append to objectBuilder"); + checkStatus(weightBuilder->Append(obj.getWeight()), "Failed to append weight"); + appendRelationRef(fromBuilder, obj.getFrom().getObjectID()); + appendRelationRef(toBuilder, obj.getTo().getObjectID()); + } + std::shared_ptr array; + checkStatus(collectionBuilder->Finish(&array), "Failed to finish collectionBuilder"); + return array; + }); {% endfor %} return true; } static const bool convertersRegistered = registerConverters(); +bool registerBufferReaders() { + auto& convRegistry = podio::ArrowConverterRegistry::mutInstance(); +{% for datatype in datatypes %} + convRegistry.registerReader("{{ datatype.class.full_type }}", [](std::shared_ptr array, int64_t rowIndex, bool isSubset, podio::SchemaVersionT version) -> std::optional { + auto buffers = podio::CollectionBufferFactory::instance().createBuffers("{{ datatype.class.full_type }}Collection", version, isSubset); + if (!buffers) { + return std::nullopt; + } + + auto collection_frames_list = std::static_pointer_cast(array); + auto frame_collection_slice = collection_frames_list->value_slice(rowIndex); + auto raw_data_struct_array = std::static_pointer_cast(frame_collection_slice); + + if (!buffers->data) { + auto collId_arr = std::static_pointer_cast(raw_data_struct_array->field(0)); + auto index_arr = std::static_pointer_cast(raw_data_struct_array->field(1)); + auto* refVec = buffers->references->at(0).get(); + size_t collection_size = raw_data_struct_array->length(); + refVec->reserve(collection_size); + for (size_t i = 0; i < collection_size; ++i) { + refVec->push_back(podio::ObjectID{index_arr->Value(i), collId_arr->Value(i)}); + } + return buffers; + } + + auto* dataVec = buffers->dataAsVector<{{ datatype.class.full_type }}Data>(); + size_t collection_size = raw_data_struct_array->length(); + dataVec->resize(collection_size); + {%- for member in datatype.arrow_metadata.members %} + {{ declare_arrays(member, 'raw_data_struct_array', '') | indent(4) }} + {%- endfor %} + {%- for member in datatype.arrow_metadata.vector_members %} + auto arr_{{ member.name }} = std::static_pointer_cast(raw_data_struct_array->GetFieldByName("{{ member.name }}")); + {%- set val_array_type = member.value_builder.arrow_array_type %} + {%- if member.value_builder.is_primitive %} + auto arr_{{ member.name }}_val = std::static_pointer_cast<{{ val_array_type }}>(arr_{{ member.name }}->values()); + {%- else %} + auto arr_{{ member.name }}_val = std::static_pointer_cast(arr_{{ member.name }}->values()); + {%- for child in member.value_builder.children %} + {{ declare_arrays(child, 'arr_' + member.name + '_val', member.name + '_val_') | indent(4) }} + {%- endfor %} + {%- endif %} + {%- endfor %} + + for (size_t i = 0; i < collection_size; ++i) { + {%- for member in datatype.arrow_metadata.members %} + {{ read_member(member, '(*dataVec)[i]', '') | indent(6) }} + {%- endfor %} + {%- for member in datatype.arrow_metadata.vector_members %} + { + auto* vecMemberBuffer = podio::CollectionReadBuffers::asVector<{{ member.full_type }}>(buffers->vectorMembers->at({{ loop.index0 }}).second); + size_t count_begin = vecMemberBuffer->size(); + // Get start and end offsets of the vector/list element in the flat values array + int32_t start_offset = arr_{{ member.name }}->value_offset(i); + int32_t end_offset = arr_{{ member.name }}->value_offset(i + 1); + for (int32_t j = start_offset; j < end_offset; ++j) { + {% if member.value_builder.is_primitive -%} + {% if member.value_builder.full_type == 'std::string' -%} + vecMemberBuffer->push_back(arr_{{ member.name }}_val->GetString(j)); + {% else -%} + vecMemberBuffer->push_back(arr_{{ member.name }}_val->Value(j)); + {% endif -%} + {% else -%} + {{ member.full_type }} val; + {% for child in member.value_builder.children -%} + {{ read_member(child, 'val', member.name + '_val_', 'j') | indent(10) }} + {% endfor -%} + vecMemberBuffer->push_back(val); + {% endif -%} + } + (*dataVec)[i].{{ member.name }}_begin = count_begin; + (*dataVec)[i].{{ member.name }}_end = vecMemberBuffer->size(); + } + {%- endfor %} + {%- for relation in datatype.arrow_metadata.one_to_many_relations %} + { + auto rel_arr = std::static_pointer_cast(raw_data_struct_array->GetFieldByName("{{ relation.name }}")); + auto rel_struct_arr = std::static_pointer_cast(rel_arr->values()); + auto rel_collId_arr = std::static_pointer_cast(rel_struct_arr->field(0)); + auto rel_index_arr = std::static_pointer_cast(rel_struct_arr->field(1)); + + auto* refVec = buffers->references->at({{ loop.index0 }}).get(); + size_t count_begin = refVec->size(); + // Get start and end offsets of the relation list in the flat values array + int32_t start_offset = rel_arr->value_offset(i); + int32_t end_offset = rel_arr->value_offset(i + 1); + for (int32_t j = start_offset; j < end_offset; ++j) { + refVec->push_back(podio::ObjectID{rel_index_arr->Value(j), rel_collId_arr->Value(j)}); + } + (*dataVec)[i].{{ relation.name }}_begin = count_begin; + (*dataVec)[i].{{ relation.name }}_end = refVec->size(); + } + {%- endfor %} + {%- for relation in datatype.arrow_metadata.one_to_one_relations %} + { + auto rel_struct_arr = std::static_pointer_cast(raw_data_struct_array->GetFieldByName("{{ relation.name }}")); + auto rel_collId_arr = std::static_pointer_cast(rel_struct_arr->field(0)); + auto rel_index_arr = std::static_pointer_cast(rel_struct_arr->field(1)); + + auto* refVec = buffers->references->at({{ loop.index0 + datatype.arrow_metadata.one_to_many_relations|length }}).get(); + refVec->push_back(podio::ObjectID{rel_index_arr->Value(i), rel_collId_arr->Value(i)}); + } + {%- endfor %} + } + return buffers; + }); +{% endfor %} +{% for link in links %} + convRegistry.registerReader("podio::Link<{{ link.From.full_type }},{{ link.To.full_type }}>", [](std::shared_ptr array, int64_t rowIndex, bool isSubset, podio::SchemaVersionT version) -> std::optional { + auto buffers = podio::CollectionBufferFactory::instance().createBuffers("podio::LinkCollection<{{ link.From.full_type }},{{ link.To.full_type }}>", version, isSubset); + if (!buffers) { + return std::nullopt; + } + auto list_array = std::static_pointer_cast(array); + auto obj_array = list_array->value_slice(rowIndex); + auto struct_array = std::static_pointer_cast(obj_array); + size_t num_objs = struct_array->length(); + + if (!buffers->data) { + auto collId_arr = std::static_pointer_cast(struct_array->field(0)); + auto index_arr = std::static_pointer_cast(struct_array->field(1)); + auto* refVec = buffers->references->at(0).get(); + refVec->reserve(num_objs); + for (size_t i = 0; i < num_objs; ++i) { + refVec->push_back(podio::ObjectID{index_arr->Value(i), collId_arr->Value(i)}); + } + return buffers; + } + + auto* dataVec = buffers->dataAsVector(); + dataVec->resize(num_objs); + + auto weight_arr = std::static_pointer_cast(struct_array->GetFieldByName("weight")); + auto from_struct_arr = std::static_pointer_cast(struct_array->GetFieldByName("from")); + auto to_struct_arr = std::static_pointer_cast(struct_array->GetFieldByName("to")); + + auto* fromRefVec = buffers->references->at(0).get(); + auto* toRefVec = buffers->references->at(1).get(); + fromRefVec->reserve(num_objs); + toRefVec->reserve(num_objs); + + auto readRelationRef = [](arrow::StructArray* structArr, size_t idx, std::vector* refVec) { + auto collId = std::static_pointer_cast(structArr->field(0))->Value(idx); + auto index = std::static_pointer_cast(structArr->field(1))->Value(idx); + refVec->push_back(podio::ObjectID{index, collId}); + }; + + for (size_t i = 0; i < num_objs; ++i) { + (*dataVec)[i].weight = weight_arr->Value(i); + readRelationRef(from_struct_arr.get(), i, fromRefVec); + readRelationRef(to_struct_arr.get(), i, toRefVec); + } + return buffers; + }); +{% endfor %} + return true; +} + +static const bool readersRegistered = registerBufferReaders(); + } // namespace } // namespace {{ package_name }}::arrow_io diff --git a/src/ArrowConverterRegistry.cc b/src/ArrowConverterRegistry.cc index 4ad92712e..7a18da609 100644 --- a/src/ArrowConverterRegistry.cc +++ b/src/ArrowConverterRegistry.cc @@ -23,4 +23,16 @@ ArrowConverterRegistry::CreatorFunc ArrowConverterRegistry::getConverter(const s return nullptr; } +void ArrowConverterRegistry::registerReader(const std::string& typeName, BufferReaderFunc&& reader) { + m_readerRegistry[typeName] = std::move(reader); +} + +ArrowConverterRegistry::BufferReaderFunc ArrowConverterRegistry::getReader(const std::string& typeName) const { + auto it = m_readerRegistry.find(typeName); + if (it != m_readerRegistry.end()) { + return it->second; + } + return nullptr; +} + } // namespace podio diff --git a/src/ArrowFrameConverter.cc b/src/ArrowFrameConverter.cc index ac34509cf..c6edbdd0f 100644 --- a/src/ArrowFrameConverter.cc +++ b/src/ArrowFrameConverter.cc @@ -1,9 +1,11 @@ #include "podio/utilities/ArrowFrameConverter.h" #include "podio/Frame.h" #include "podio/utilities/ArrowConverterRegistry.h" +#include "podio/utilities/ArrowFrameData.h" #include "podio/utilities/ArrowTypeRegistry.h" #include "podio/utilities/ArrowUtils.h" +#include #include #include @@ -52,10 +54,10 @@ namespace { auto stringMap = buildParamMap(params); std::vector> fields = { - arrow::field("int_params", arrow::map(arrow::utf8(), arrow::list(arrow::int32())), true), - arrow::field("float_params", arrow::map(arrow::utf8(), arrow::list(arrow::float32())), true), - arrow::field("double_params", arrow::map(arrow::utf8(), arrow::list(arrow::float64())), true), - arrow::field("string_params", arrow::map(arrow::utf8(), arrow::list(arrow::utf8())), true), + arrow::field("int_params", arrow::map(arrow::utf8(), arrow::list(arrow::int32()))), + arrow::field("float_params", arrow::map(arrow::utf8(), arrow::list(arrow::float32()))), + arrow::field("double_params", arrow::map(arrow::utf8(), arrow::list(arrow::float64()))), + arrow::field("string_params", arrow::map(arrow::utf8(), arrow::list(arrow::utf8()))), }; auto structResult = arrow::StructArray::Make({intMap, floatMap, doubleMap, stringMap}, fields); @@ -144,10 +146,11 @@ std::shared_ptr convertFrameToTable(const podio::Frame& frame, std::to_string(array->length()) + ", expected 1"); } - // Attach "value_type", "is_subset", and "coll_id" metadata to the field - auto metadata = arrow::KeyValueMetadata::Make( - {"value_type", "is_subset", "coll_id"}, - {typeName, coll->isSubsetCollection() ? "1" : "0", std::to_string(coll->getID())}); + // Attach "value_type", "is_subset", "coll_id", and "schema_version" metadata to the field + auto metadata = + arrow::KeyValueMetadata::Make({"value_type", "is_subset", "coll_id", "schema_version"}, + {typeName, coll->isSubsetCollection() ? "1" : "0", std::to_string(coll->getID()), + std::to_string(coll->getSchemaVersion())}); auto field = arrow::field(collName, arrowType, /*nullable=*/true, std::move(metadata)); fields.push_back(std::move(field)); @@ -176,4 +179,8 @@ std::shared_ptr convertFrameToTable(const podio::Frame& frame, return tableResult.ValueOrDie(); } +podio::Frame convertTableToFrame(const std::shared_ptr& table, int rowIndex) { + return podio::Frame(std::make_unique(table, rowIndex)); +} + } // namespace podio diff --git a/src/ArrowFrameData.cc b/src/ArrowFrameData.cc new file mode 100644 index 000000000..a24bbfcb4 --- /dev/null +++ b/src/ArrowFrameData.cc @@ -0,0 +1,204 @@ +#include "podio/utilities/ArrowFrameData.h" +#include "podio/utilities/ArrowConverterRegistry.h" + +#include +#include +#include +#include + +namespace podio { + +namespace { + + struct IntSetter { + using ArrayType = arrow::Int32Array; + using ValueType = int32_t; + static ValueType getValue(const ArrayType* arr, int32_t idx) { + return arr->Value(idx); + } + static void set(podio::GenericParameters* p, const std::string& k, const std::vector& v) { + p->set(k, v); + } + }; + struct FloatSetter { + using ArrayType = arrow::FloatArray; + using ValueType = float; + static ValueType getValue(const ArrayType* arr, int32_t idx) { + return arr->Value(idx); + } + static void set(podio::GenericParameters* p, const std::string& k, const std::vector& v) { + p->set(k, v); + } + }; + struct DoubleSetter { + using ArrayType = arrow::DoubleArray; + using ValueType = double; + static ValueType getValue(const ArrayType* arr, int32_t idx) { + return arr->Value(idx); + } + static void set(podio::GenericParameters* p, const std::string& k, const std::vector& v) { + p->set(k, v); + } + }; + struct StringSetter { + using ArrayType = arrow::StringArray; + using ValueType = std::string; + static ValueType getValue(const ArrayType* arr, int32_t idx) { + return arr->GetString(idx); + } + static void set(podio::GenericParameters* p, const std::string& k, const std::vector& v) { + p->set(k, v); + } + }; + + template + void extractMap(const arrow::StructArray* struct_array, const std::string& fieldName, int64_t rowIndex, + podio::GenericParameters* params) { + auto map_array = std::static_pointer_cast(struct_array->GetFieldByName(fieldName)); + if (!map_array) { + return; + } + auto keys_array = std::static_pointer_cast(map_array->keys()); + auto items_array = map_array->items(); + auto list_items = std::static_pointer_cast(items_array); + + int32_t start = map_array->value_offset(rowIndex); + int32_t end = map_array->value_offset(rowIndex + 1); + + using ArrayType = typename Setter::ArrayType; + using ValueType = typename Setter::ValueType; + auto val_array = std::static_pointer_cast(list_items->values()); + + for (int32_t i = start; i < end; ++i) { + std::string key = keys_array->GetString(i); + + int32_t list_start = list_items->value_offset(i); + int32_t list_end = list_items->value_offset(i + 1); + + std::vector vals; + vals.reserve(list_end - list_start); + for (int32_t j = list_start; j < list_end; ++j) { + vals.push_back(Setter::getValue(val_array.get(), j)); + } + Setter::set(params, key, vals); + } + } + + std::shared_ptr getChunkOrThrow(const std::shared_ptr& chunked_array, + int64_t rowIndex, int64_t& remaining_idx) { + remaining_idx = rowIndex; + for (const auto& c : chunked_array->chunks()) { + if (remaining_idx < c->length()) { + return c; + } + remaining_idx -= c->length(); + } + throw std::runtime_error("Row index out of bounds in chunked array"); + } + +} // namespace + +ArrowFrameData::ArrowFrameData(std::shared_ptr table, int64_t rowIndex) : + m_table(std::move(table)), m_rowIndex(rowIndex), m_availableCollections(), m_idTable() { + if (!m_table) { + throw std::runtime_error("ArrowTable is null"); + } + if (m_rowIndex < 0 || m_rowIndex >= m_table->num_rows()) { + throw std::runtime_error("ArrowTable row index out of bounds"); + } + + std::vector ids; + std::vector names; + + 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()); + + auto metadata = field->metadata(); + if (!metadata) { + throw std::runtime_error("Missing metadata for collection column: " + field->name()); + } + auto collIdIdx = metadata->FindKey("coll_id"); + if (collIdIdx == -1) { + throw std::runtime_error("Missing coll_id in metadata for collection: " + field->name()); + } + + uint32_t collId = std::stoul(metadata->value(collIdIdx)); + if (std::find(ids.begin(), ids.end(), collId) != ids.end() || + std::find(names.begin(), names.end(), field->name()) != names.end()) { + throw std::runtime_error("Duplicate collection ID or name: " + field->name()); + } + ids.push_back(collId); + names.push_back(field->name()); + } + m_idTable = podio::CollectionIDTable(std::move(ids), std::move(names)); +} + +std::optional ArrowFrameData::getCollectionBuffers(const std::string& name) { + auto chunked_array = m_table->GetColumnByName(name); + if (!chunked_array) { + return std::nullopt; + } + + auto field = m_table->schema()->GetFieldByName(name); + auto metadata = field->metadata(); + + auto typeNameIdx = metadata->FindKey("value_type"); + auto isSubsetIdx = metadata->FindKey("is_subset"); + + if (typeNameIdx == -1 || isSubsetIdx == -1) { + throw std::runtime_error("Missing value_type or is_subset in metadata for: " + name); + } + + std::string typeName = metadata->value(typeNameIdx); + bool isSubset = (metadata->value(isSubsetIdx) == "1"); + + auto schemaVersionIdx = metadata->FindKey("schema_version"); + uint32_t schemaVersion = 1; + if (schemaVersionIdx != -1) { + schemaVersion = std::stoul(metadata->value(schemaVersionIdx)); + } + + int64_t remaining_idx = 0; + auto chunk = getChunkOrThrow(chunked_array, m_rowIndex, remaining_idx); + + auto reader = podio::ArrowConverterRegistry::instance().getReader(typeName); + if (!reader) { + throw std::runtime_error("No Arrow reader registered for type: " + typeName); + } + + return reader(chunk, remaining_idx, isSubset, schemaVersion); +} + +std::unique_ptr ArrowFrameData::getParameters() { + auto chunked_array = m_table->GetColumnByName("frame_parameters"); + if (!chunked_array) { + return std::make_unique(); + } + int64_t remaining_idx = 0; + auto chunk = getChunkOrThrow(chunked_array, m_rowIndex, remaining_idx); + + auto struct_array = std::static_pointer_cast(chunk); + auto params = std::make_unique(); + + extractMap(struct_array.get(), "int_params", remaining_idx, params.get()); + extractMap(struct_array.get(), "float_params", remaining_idx, params.get()); + extractMap(struct_array.get(), "double_params", remaining_idx, params.get()); + extractMap(struct_array.get(), "string_params", remaining_idx, params.get()); + + return params; +} + +std::vector ArrowFrameData::getAvailableCollections() const { + return m_availableCollections; +} + +podio::CollectionIDTable ArrowFrameData::getIDTable() const { + return {m_idTable.ids(), m_idTable.names()}; +} + +} // namespace podio diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index a47f5032d..759a145af 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -156,6 +156,7 @@ if(ENABLE_ARROW) ArrowUtils.cc ArrowTypeRegistry.cc ArrowConverterRegistry.cc + ArrowFrameData.cc ) add_library(podioArrow SHARED ${arrow_sources}) diff --git a/tests/unittests/test_arrow_converter.cpp b/tests/unittests/test_arrow_converter.cpp index 995adb819..a0ccf0678 100644 --- a/tests/unittests/test_arrow_converter.cpp +++ b/tests/unittests/test_arrow_converter.cpp @@ -13,62 +13,22 @@ // Arrow Frame Converter #include "podio/utilities/ArrowFrameConverter.h" +#include "podio/utilities/ArrowTypeRegistry.h" + +// Test frame validation helpers +#include "read_frame.h" +#include "write_frame.h" // Arrow headers #include +#include #include #include -#include TEST_CASE("ArrowFrameConverter - convertFrameToTable Verification", "[arrow][converter]") { - podio::Frame originalFrame; - - // 1. Prepare some hits - ExampleHitCollection hits; - auto hit1 = MutableExampleHit(0x42ULL, 1.0f, 2.0f, 3.0f, 10.5); - auto hit2 = MutableExampleHit(0x43ULL, 4.0f, 5.0f, 6.0f, 20.5); - hits.push_back(hit1); - hits.push_back(hit2); - - // 2. Prepare some clusters with relations to hits - ExampleClusterCollection clusters; - auto cluster1 = MutableExampleCluster(); - cluster1.energy(100.0); - cluster1.addHits(hit1); - cluster1.addHits(hit2); - clusters.push_back(cluster1); - - // 3. Prepare a relation collection - ExampleWithOneRelationCollection relColls; - auto relObj = MutableExampleWithOneRelation(); - relObj.cluster(cluster1); - relColls.push_back(relObj); - - // 4. Prepare a vector member collection - ExampleWithVectorMemberCollection vecColls; - auto vecObj = MutableExampleWithVectorMember(); - vecObj.addcount(42); - vecObj.addcount(137); - vecColls.push_back(vecObj); - - // 5. Prepare a subset collection of hits - ExampleHitCollection subsetHits; - subsetHits.setSubsetCollection(true); - subsetHits.push_back(hit1); - - // Put them all into the frame - originalFrame.put(std::move(hits), "Hits"); - originalFrame.put(std::move(clusters), "Clusters"); - originalFrame.put(std::move(relColls), "OneRelation"); - originalFrame.put(std::move(vecColls), "VectorMember"); - originalFrame.put(std::move(subsetHits), "SubsetHits"); - - // Put some frame parameters - originalFrame.putParameter("anInt", 42); - originalFrame.putParameter("someFloats", std::vector{1.23f, 2.34f, 3.45f}); - originalFrame.putParameter("someStrings", std::vector{"hello", "world"}); - - std::vector colls = {"Hits", "Clusters", "OneRelation", "VectorMember", "SubsetHits"}; + auto originalFrame = makeFrame(0); + + std::vector colls = {"hits", "clusters", "OneRelation", "WithVectorMember", "hitRefs"}; // Convert to Arrow Table auto table = podio::convertFrameToTable(originalFrame, colls); @@ -81,83 +41,86 @@ TEST_CASE("ArrowFrameConverter - convertFrameToTable Verification", "[arrow][con // Check columns present auto schema = table->schema(); - REQUIRE(schema->GetFieldByName("Hits") != nullptr); - REQUIRE(schema->GetFieldByName("Clusters") != nullptr); + REQUIRE(schema->GetFieldByName("hits") != nullptr); + REQUIRE(schema->GetFieldByName("clusters") != nullptr); REQUIRE(schema->GetFieldByName("OneRelation") != nullptr); - REQUIRE(schema->GetFieldByName("VectorMember") != nullptr); - REQUIRE(schema->GetFieldByName("SubsetHits") != nullptr); + REQUIRE(schema->GetFieldByName("WithVectorMember") != nullptr); + REQUIRE(schema->GetFieldByName("hitRefs") != nullptr); REQUIRE(schema->GetFieldByName("frame_parameters") != nullptr); - // Assertions on Hits table structure and metadata - auto hitsField = schema->GetFieldByName("Hits"); + // Assertions on hits table structure and metadata + auto hitsField = schema->GetFieldByName("hits"); auto hitsMeta = hitsField->metadata(); REQUIRE(hitsMeta != nullptr); REQUIRE(hitsMeta->Get("value_type").ValueOrDie() == "ExampleHit"); REQUIRE(hitsMeta->Get("is_subset").ValueOrDie() == "0"); - auto hitsArray = std::static_pointer_cast(table->GetColumnByName("Hits")->chunk(0)); + auto hitsArray = std::static_pointer_cast(table->GetColumnByName("hits")->chunk(0)); REQUIRE(hitsArray != nullptr); REQUIRE(hitsArray->length() == 1); REQUIRE(hitsArray->value_length(0) == 2); auto hitsStruct = std::static_pointer_cast(hitsArray->values()); auto hitsX = std::static_pointer_cast(hitsStruct->GetFieldByName("x")); - REQUIRE(hitsX->Value(0) == 1.0); - REQUIRE(hitsX->Value(1) == 4.0); + REQUIRE(hitsX->Value(0) == 0.0); + REQUIRE(hitsX->Value(1) == 1.0); auto hitsEnergy = std::static_pointer_cast(hitsStruct->GetFieldByName("energy")); - REQUIRE(hitsEnergy->Value(0) == 10.5); - REQUIRE(hitsEnergy->Value(1) == 20.5); + REQUIRE(hitsEnergy->Value(0) == 23.0); + REQUIRE(hitsEnergy->Value(1) == 12.0); - // Assertions on Clusters table structure and metadata - auto clustersField = schema->GetFieldByName("Clusters"); + // Assertions on clusters table structure and metadata + auto clustersField = schema->GetFieldByName("clusters"); auto clustersMeta = clustersField->metadata(); REQUIRE(clustersMeta->Get("value_type").ValueOrDie() == "ExampleCluster"); - auto clustersArray = std::static_pointer_cast(table->GetColumnByName("Clusters")->chunk(0)); - REQUIRE(clustersArray->value_length(0) == 1); + auto clustersArray = std::static_pointer_cast(table->GetColumnByName("clusters")->chunk(0)); + REQUIRE(clustersArray->value_length(0) == 3); auto clustersStruct = std::static_pointer_cast(clustersArray->values()); auto clustersEnergy = std::static_pointer_cast(clustersStruct->GetFieldByName("energy")); - REQUIRE(clustersEnergy->Value(0) == 100.0); + REQUIRE(clustersEnergy->Value(0) == 23.0); + REQUIRE(clustersEnergy->Value(1) == 12.0); + REQUIRE(clustersEnergy->Value(2) == 35.0); auto clustersHits = std::static_pointer_cast(clustersStruct->GetFieldByName("Hits")); - REQUIRE(clustersHits->value_length(0) == 2); + REQUIRE(clustersHits->value_length(0) == 1); + REQUIRE(clustersHits->value_length(2) == 2); auto clustersHitsStruct = std::static_pointer_cast(clustersHits->values()); auto clustersHitsIndex = std::static_pointer_cast(clustersHitsStruct->GetFieldByName("index")); - REQUIRE(clustersHitsIndex->Value(0) == 0); // hit1 index - REQUIRE(clustersHitsIndex->Value(1) == 1); // hit2 index + REQUIRE(clustersHitsIndex->Value(0) == 0); // hit0 index // Assertions on OneRelation auto relArray = std::static_pointer_cast(table->GetColumnByName("OneRelation")->chunk(0)); - REQUIRE(relArray->value_length(0) == 1); + REQUIRE(relArray->value_length(0) == 2); auto relStruct = std::static_pointer_cast(relArray->values()); auto relCluster = std::static_pointer_cast(relStruct->GetFieldByName("cluster")); auto relClusterIndex = std::static_pointer_cast(relCluster->GetFieldByName("index")); - REQUIRE(relClusterIndex->Value(0) == 0); // cluster1 index + REQUIRE(relClusterIndex->Value(0) == 2); // cluster index 2 - // Assertions on VectorMember - auto vecArray = std::static_pointer_cast(table->GetColumnByName("VectorMember")->chunk(0)); - REQUIRE(vecArray->value_length(0) == 1); + // Assertions on WithVectorMember + auto vecArray = std::static_pointer_cast(table->GetColumnByName("WithVectorMember")->chunk(0)); + REQUIRE(vecArray->value_length(0) == 2); auto vecStruct = std::static_pointer_cast(vecArray->values()); auto vecCount = std::static_pointer_cast(vecStruct->GetFieldByName("count")); auto vecCountVal = std::static_pointer_cast(vecCount->values()); - REQUIRE(vecCountVal->Value(0) == 42); - REQUIRE(vecCountVal->Value(1) == 137); + REQUIRE(vecCountVal->Value(0) == 0); + REQUIRE(vecCountVal->Value(1) == 10); - // Assertions on SubsetHits - auto subField = schema->GetFieldByName("SubsetHits"); + // Assertions on hitRefs (subset collection) + auto subField = schema->GetFieldByName("hitRefs"); REQUIRE(subField->metadata()->Get("is_subset").ValueOrDie() == "1"); - auto subArray = std::static_pointer_cast(table->GetColumnByName("SubsetHits")->chunk(0)); - REQUIRE(subArray->value_length(0) == 1); + auto subArray = std::static_pointer_cast(table->GetColumnByName("hitRefs")->chunk(0)); + REQUIRE(subArray->value_length(0) == 2); auto subObjectID = std::static_pointer_cast(subArray->values()); auto subIndex = std::static_pointer_cast(subObjectID->GetFieldByName("index")); - REQUIRE(subIndex->Value(0) == 0); // pointing to hit1 index + REQUIRE(subIndex->Value(0) == 1); // pointing to hits[1] + REQUIRE(subIndex->Value(1) == 0); // pointing to hits[0] // Assertions on frame parameters auto paramArray = std::static_pointer_cast(table->GetColumnByName("frame_parameters")->chunk(0)); @@ -168,314 +131,349 @@ TEST_CASE("ArrowFrameConverter - convertFrameToTable Verification", "[arrow][con auto intParams = std::static_pointer_cast(paramArray->GetFieldByName("int_params")); REQUIRE(intParams != nullptr); REQUIRE(intParams->length() == 1); + REQUIRE(intParams->value_length(0) == 3); auto intKeys = std::static_pointer_cast(intParams->keys()); - auto intItems = std::static_pointer_cast(intParams->items()); - REQUIRE(intKeys->GetString(0) == "anInt"); - auto intValues = std::static_pointer_cast(intItems->values()); - REQUIRE(intValues->Value(0) == 42); - - // Verify float_params - auto floatParams = std::static_pointer_cast(paramArray->GetFieldByName("float_params")); - REQUIRE(floatParams != nullptr); - REQUIRE(floatParams->length() == 1); - auto floatKeys = std::static_pointer_cast(floatParams->keys()); - auto floatItems = std::static_pointer_cast(floatParams->items()); - REQUIRE(floatKeys->GetString(0) == "someFloats"); - auto floatValues = std::static_pointer_cast(floatItems->values()); - REQUIRE(floatValues->Value(0) == 1.23f); - REQUIRE(floatValues->Value(1) == 2.34f); - REQUIRE(floatValues->Value(2) == 3.45f); - - // Verify string_params - auto stringParams = std::static_pointer_cast(paramArray->GetFieldByName("string_params")); - REQUIRE(stringParams != nullptr); - REQUIRE(stringParams->length() == 1); - auto stringKeys = std::static_pointer_cast(stringParams->keys()); - auto stringItems = std::static_pointer_cast(stringParams->items()); - REQUIRE(stringKeys->GetString(0) == "someStrings"); - auto stringValues = std::static_pointer_cast(stringItems->values()); - REQUIRE(stringValues->GetString(0) == "hello"); - REQUIRE(stringValues->GetString(1) == "world"); + REQUIRE(intKeys->GetString(0) == "SomeValue"); + REQUIRE(intKeys->GetString(1) == "SomeVectorData"); + REQUIRE(intKeys->GetString(2) == "anInt"); } -// ============================================================================ -// Arrow Table & Frame Parameters JSON Serialization Helpers -// ============================================================================ - -namespace { - -// wrap_relation: -// In PODIO's native JSON output: -// - A one-to-one relation is wrapped in a 1-element array: "cluster": [{"collectionID": X, "index": Y}] -// - A one-to-many relation or subset collection is a simple flat list of elements: "Hits": [{"collectionID": X, -// "index": Y}, ...] -// Since the Arrow Table structures them both as relation reference structs, we use wrap_relation=false inside list -// loops (like listToJson) to prevent wrapping each individual element of a one-to-many list in its own array. -nlohmann::json arrayElementToJson(const arrow::Array& array, int64_t idx, bool wrap_relation = true); - -nlohmann::json structToJson(const arrow::StructArray& struct_array, int64_t idx, bool wrap_relation = true) { - auto struct_type = struct_array.struct_type(); - - bool is_relation = false; - if (struct_array.num_fields() == 2) { - std::string f0 = struct_type->field(0)->name(); - std::string f1 = struct_type->field(1)->name(); - if ((f0 == "collectionID" && f1 == "index") || (f0 == "index" && f1 == "collectionID")) { - is_relation = true; - } +TEST_CASE("ArrowFrameConverter - convertTableToFrame Reader-Only Verification", "[arrow][converter][reader]") { + auto arrowType = podio::ArrowTypeRegistry::instance().getType("ExampleHit"); + REQUIRE(arrowType != nullptr); + + std::unique_ptr builder; + auto status = arrow::MakeBuilder(arrow::default_memory_pool(), arrowType, &builder); + REQUIRE(status.ok()); + auto* collectionBuilder = static_cast(builder.get()); + auto* objectBuilder = static_cast(collectionBuilder->value_builder()); + auto* builder_cellID = static_cast(objectBuilder->child(0)); + auto* builder_x = static_cast(objectBuilder->child(1)); + auto* builder_y = static_cast(objectBuilder->child(2)); + auto* builder_z = static_cast(objectBuilder->child(3)); + auto* builder_energy = static_cast(objectBuilder->child(4)); + + REQUIRE(collectionBuilder->Append().ok()); + + // Hit 1 + REQUIRE(builder_cellID->Append(0x100ULL).ok()); + REQUIRE(builder_x->Append(10.0).ok()); + REQUIRE(builder_y->Append(20.0).ok()); + REQUIRE(builder_z->Append(30.0).ok()); + REQUIRE(builder_energy->Append(5.5).ok()); + REQUIRE(objectBuilder->Append().ok()); + + // Hit 2 + REQUIRE(builder_cellID->Append(0x200ULL).ok()); + REQUIRE(builder_x->Append(-1.0).ok()); + REQUIRE(builder_y->Append(-2.0).ok()); + REQUIRE(builder_z->Append(-3.0).ok()); + REQUIRE(builder_energy->Append(0.5).ok()); + REQUIRE(objectBuilder->Append().ok()); + + std::shared_ptr array; + REQUIRE(collectionBuilder->Finish(&array).ok()); + + auto metadata = arrow::KeyValueMetadata::Make({"value_type", "is_subset", "coll_id"}, {"ExampleHit", "0", "101"}); + auto field = arrow::field("Hits", arrowType, true, std::move(metadata)); + + auto schema = arrow::schema({field}); + auto batch = arrow::RecordBatch::Make(schema, 1, {array}); + REQUIRE(batch != nullptr); + auto tableResult = arrow::Table::FromRecordBatches({batch}); + REQUIRE(tableResult.ok()); + const auto& table = tableResult.ValueOrDie(); + + auto frame = podio::convertTableToFrame(table, 0); + + const auto& recHits = frame.get("Hits"); + REQUIRE(recHits.size() == 2); + REQUIRE(recHits[0].cellID() == 0x100ULL); + REQUIRE(recHits[0].x() == 10.0f); + REQUIRE(recHits[0].energy() == 5.5); + REQUIRE(recHits[1].cellID() == 0x200ULL); + REQUIRE(recHits[1].x() == -1.0f); + REQUIRE(recHits[1].energy() == 0.5); +} + +void verifyEventNoUserData(const podio::Frame& event, int eventNum) { + const float evtWeight = event.getParameter("UserEventWeight").value(); + REQUIRE(evtWeight == 100.f * eventNum); + + std::stringstream ss; + ss << " event_number_" << eventNum; + const auto evtName = event.getParameter("UserEventName").value(); + REQUIRE(evtName == ss.str()); + + const auto someVectorData = event.getParameter>("SomeVectorData").value(); + REQUIRE(someVectorData.size() == 4); + for (int i = 0; i < 4; ++i) { + REQUIRE(someVectorData[i] == i + 1); } - if (is_relation) { - auto ref = nlohmann::json::object(); - for (int f = 0; f < struct_array.num_fields(); ++f) { - std::string name = struct_type->field(f)->name(); - const auto& field_array = struct_array.field(f); - ref[name] = arrayElementToJson(*field_array, idx, true); + const auto doubleParams = event.getParameter>("SomeVectorData").value(); + REQUIRE(doubleParams.size() == 2); + REQUIRE(doubleParams[0] == eventNum * 1.1); + REQUIRE(doubleParams[1] == eventNum * 2.2); + + checkHitCollection(event, eventNum); + + auto& hits = event.get("hits"); + + auto& hitRefs = event.get("hitRefs"); + REQUIRE(hitRefs.size() == hits.size()); + REQUIRE((hits[1] == hitRefs[0] && hits[0] == hitRefs[1])); + + checkClusterCollection(event, hits); + + auto& mcpRefs = event.get("mcParticleRefs"); + for (auto ref : mcpRefs) { + const auto daughters = ref.daughters(); + if (!daughters.empty()) { + REQUIRE(daughters[0].isAvailable()); } - if (wrap_relation) { - return nlohmann::json::array({ref}); - } else { - return ref; + const auto parents = ref.parents(); + if (!parents.empty()) { + REQUIRE(parents[0].isAvailable()); } } - auto j = nlohmann::json::object(); - for (int f = 0; f < struct_array.num_fields(); ++f) { - std::string name = struct_type->field(f)->name(); - const auto& field_array = struct_array.field(f); - j[name] = arrayElementToJson(*field_array, idx, true); - } - return j; -} + checkMCParticleCollection(event, podio::version::build_version); -nlohmann::json listToJson(const arrow::ListArray& list_array, int64_t idx) { - auto j = nlohmann::json::array(); - int64_t offset = list_array.value_offset(idx); - int64_t length = list_array.value_length(idx); - const auto& values = list_array.values(); - for (int64_t val_i = 0; val_i < length; ++val_i) { - j.push_back(arrayElementToJson(*values, offset + val_i, false)); - } - return j; -} + auto& mcps = event.get("mcparticles"); -nlohmann::json fixedSizeListToJson(const arrow::FixedSizeListArray& list_array, int64_t idx) { - auto j = nlohmann::json::array(); - int64_t offset = list_array.value_offset(idx); - int64_t length = list_array.value_length(); - const auto& values = list_array.values(); - for (int64_t val_i = 0; val_i < length; ++val_i) { - j.push_back(arrayElementToJson(*values, offset + val_i, false)); + for (auto pr : mcpRefs) { + REQUIRE(static_cast(pr.getObjectID().collectionID) != mcpRefs.getID()); } - return j; -} -nlohmann::json mapToJson(const arrow::MapArray& map_array, int64_t idx) { - auto j = nlohmann::json::object(); - int64_t offset = map_array.value_offset(idx); - int64_t length = map_array.value_length(idx); - const auto& keys = map_array.keys(); - const auto& items = map_array.items(); - - auto string_keys = std::static_pointer_cast(keys); - for (int64_t k = 0; k < length; ++k) { - std::string key = string_keys->GetString(offset + k); - j[key] = arrayElementToJson(*items, offset + k, true); - } - return j; -} + auto& moreMCs = event.get("moreMCs"); + REQUIRE(mcps.size() == moreMCs.size()); -nlohmann::json arrayElementToJson(const arrow::Array& array, int64_t idx, bool wrap_relation) { - if (array.IsNull(idx)) { - return nullptr; + for (size_t index = 0; index < mcps.size(); ++index) { + REQUIRE(mcps[index].energy() == moreMCs[index].energy()); + REQUIRE(mcps[index].daughters().size() == moreMCs[index].daughters().size()); } - switch (array.type()->id()) { - case arrow::Type::BOOL: - return static_cast(array).Value(idx); - case arrow::Type::INT8: - return static_cast(array).Value(idx); - case arrow::Type::UINT8: - return static_cast(array).Value(idx); - case arrow::Type::INT16: - return static_cast(array).Value(idx); - case arrow::Type::UINT16: - return static_cast(array).Value(idx); - case arrow::Type::INT32: - return static_cast(array).Value(idx); - case arrow::Type::UINT32: - return static_cast(array).Value(idx); - case arrow::Type::INT64: - return static_cast(array).Value(idx); - case arrow::Type::UINT64: - return static_cast(array).Value(idx); - case arrow::Type::FLOAT: - return static_cast(array).Value(idx); - case arrow::Type::DOUBLE: - return static_cast(array).Value(idx); - case arrow::Type::STRING: - return static_cast(array).GetString(idx); - case arrow::Type::STRUCT: - return structToJson(static_cast(array), idx, wrap_relation); - case arrow::Type::LIST: - return listToJson(static_cast(array), idx); - case arrow::Type::FIXED_SIZE_LIST: - return fixedSizeListToJson(static_cast(array), idx); - case arrow::Type::MAP: - return mapToJson(static_cast(array), idx); - default: - throw std::runtime_error("Unsupported Arrow type in JSON converter: " + array.type()->ToString()); - } -} -nlohmann::json arrowTableToJson(const std::shared_ptr& table) { - auto j = nlohmann::json::object(); - for (int c = 0; c < table->num_columns(); ++c) { - std::string name = table->schema()->field(c)->name(); - auto chunked_arr = table->column(c); - auto chunk = chunked_arr->chunk(0); - j[name] = arrayElementToJson(*chunk, 0); + REQUIRE(mcpRefs.size() == mcps.size()); + for (size_t i = 0; i < mcpRefs.size(); ++i) { + if (i < 5) { + REQUIRE(mcpRefs[i] == mcps[2 * i + 1]); + } else { + const int index = (i - 5) * 2; + REQUIRE(mcpRefs[i] == moreMCs[index]); + } } - return j; -} -nlohmann::json genericParametersToJson(const podio::GenericParameters& params) { - auto j = nlohmann::json::object(); + const auto& rels = event.get("OneRelation"); + REQUIRE(rels.size() == 2); - auto int_params = nlohmann::json::object(); - for (const auto& key : params.getKeys()) { - int_params[key] = params.get>(key).value_or(std::vector{}); - } - j["int_params"] = int_params; + auto& vecs = event.get("WithVectorMember"); + REQUIRE(vecs.size() == 2); - auto float_params = nlohmann::json::object(); - for (const auto& key : params.getKeys()) { - float_params[key] = params.get>(key).value_or(std::vector{}); + for (auto vec : vecs) { + REQUIRE(vec.count().size() == 2); } - j["float_params"] = float_params; - - auto double_params = nlohmann::json::object(); - for (const auto& key : params.getKeys()) { - double_params[key] = params.get>(key).value_or(std::vector{}); - } - j["double_params"] = double_params; - - auto string_params = nlohmann::json::object(); - for (const auto& key : params.getKeys()) { - string_params[key] = params.get>(key).value_or(std::vector{}); + REQUIRE(vecs[0].count(0) == eventNum); + REQUIRE(vecs[0].count(1) == eventNum + 10); + REQUIRE(vecs[1].count(0) == eventNum + 1); + REQUIRE(vecs[1].count(1) == eventNum + 11); + + auto& arrays = event.get("arrays"); + REQUIRE(!arrays.empty()); + auto array = arrays[0]; + REQUIRE(array.myArray(1) == eventNum); + REQUIRE(array.arrayStruct().data.p.at(2) == 2 * eventNum); + REQUIRE(array.structArray(0).x == eventNum); + + auto& nmspaces = event.get("WithNamespaceRelation"); + auto& copies = event.get("WithNamespaceRelationCopy"); + + auto cpytest = ex42::ExampleWithARelationCollection{}; + for (size_t j = 0; j < nmspaces.size(); j++) { + auto nmsp = nmspaces[j]; + auto cpy = copies[j]; + cpytest.push_back(nmsp.clone()); + if (nmsp.ref().isAvailable()) { + REQUIRE(nmsp.ref().component().x == cpy.ref().component().x); + REQUIRE(nmsp.ref().component().y == cpy.ref().component().y); + REQUIRE(nmsp.ref().x() == cpy.ref().x()); + REQUIRE(nmsp.number() == cpy.number()); + REQUIRE(nmsp.ref().getObjectID() == cpy.ref().getObjectID()); + } + auto cpy_it = cpy.refs_begin(); + for (auto it = nmsp.refs_begin(); it != nmsp.refs_end(); ++it, ++cpy_it) { + REQUIRE(it->component().x == cpy_it->component().x); + REQUIRE(it->component().y == cpy_it->component().y); + REQUIRE(it->getObjectID() == cpy_it->getObjectID()); + } } - j["string_params"] = string_params; - return j; + const auto& fixedWidthInts = event.get("fixedWidthInts"); + REQUIRE(fixedWidthInts.size() == 3); + + auto maxValues = fixedWidthInts[0]; + const auto& maxComps = maxValues.fixedWidthStruct(); + REQUIRE(maxValues.fixedI16() == std::numeric_limits::max()); + REQUIRE(maxValues.fixedU32() == std::numeric_limits::max()); + REQUIRE(maxValues.fixedU64() == std::numeric_limits::max()); + REQUIRE(maxComps.fixedInteger64 == std::numeric_limits::max()); + REQUIRE(maxComps.fixedInteger32 == std::numeric_limits::max()); + REQUIRE(maxComps.fixedUnsigned16 == std::numeric_limits::max()); + + auto minValues = fixedWidthInts[1]; + const auto& minComps = minValues.fixedWidthStruct(); + REQUIRE(minValues.fixedI16() == std::numeric_limits::min()); + REQUIRE(minValues.fixedU32() == std::numeric_limits::min()); + REQUIRE(minValues.fixedU64() == std::numeric_limits::min()); + REQUIRE(minComps.fixedInteger64 == std::numeric_limits::min()); + REQUIRE(minComps.fixedInteger32 == std::numeric_limits::min()); + REQUIRE(minComps.fixedUnsigned16 == std::numeric_limits::min()); + + auto arbValues = fixedWidthInts[2]; + const auto& arbComps = arbValues.fixedWidthStruct(); + REQUIRE(arbValues.fixedI16() == std::int16_t{-12345}); + REQUIRE(arbValues.fixedU32() == std::uint32_t{1234567890}); + REQUIRE(arbValues.fixedU64() == std::uint64_t{1234567890123456789}); + REQUIRE(arbComps.fixedInteger64 == std::int64_t{-1234567890123456789}); + REQUIRE(arbComps.fixedInteger32 == std::int64_t{-1234567890}); + REQUIRE(arbComps.fixedUnsigned16 == std::uint16_t{12345}); } -} // namespace - -TEST_CASE("ArrowFrameConverter - JSON Equivalence Check", "[arrow][converter][json]") { - podio::Frame originalFrame; - - // 1. Prepare some hits - ExampleHitCollection hits; - auto hit1 = MutableExampleHit(0x42ULL, 1.0f, 2.0f, 3.0f, 10.5); - auto hit2 = MutableExampleHit(0x43ULL, 4.0f, 5.0f, 6.0f, 20.5); - hits.push_back(hit1); - hits.push_back(hit2); - - // 2. Prepare some clusters with relations to hits - ExampleClusterCollection clusters; - auto cluster1 = MutableExampleCluster(); - cluster1.energy(100.0); - cluster1.addHits(hit1); - cluster1.addHits(hit2); - clusters.push_back(cluster1); - - // 3. Prepare a relation collection - ExampleWithOneRelationCollection relColls; - auto relObj = MutableExampleWithOneRelation(); - relObj.cluster(cluster1); - relColls.push_back(relObj); - // Add an unset relation reference object - auto unsetRelObj = MutableExampleWithOneRelation(); - relColls.push_back(unsetRelObj); - - // 4. Prepare a vector member collection - ExampleWithVectorMemberCollection vecColls; - auto vecObj = MutableExampleWithVectorMember(); - vecObj.addcount(42); - vecObj.addcount(137); - vecColls.push_back(vecObj); - - // 5. Prepare a subset collection of hits - ExampleHitCollection subsetHits; - subsetHits.setSubsetCollection(true); - subsetHits.push_back(hit1); - - // 6. Prepare a collection with fixed-size arrays (FixedSizeList) - ExampleWithArrayCollection arrays; - std::array arrayTest = {0, 1, 2, 3}; - std::array arrayTest2 = {4, 5, 6, 7}; - NotSoSimpleStruct a; - a.data.p = arrayTest2; - ex2::NamespaceStruct nstruct; - nstruct.x = 42; - std::array structArrayTest = {nstruct, nstruct, nstruct, nstruct}; - auto arrayObj = MutableExampleWithArray(a, arrayTest, arrayTest, arrayTest, arrayTest, structArrayTest); - arrays.push_back(arrayObj); - - // 7. Prepare an empty collection - ExampleClusterCollection emptyClusterColl; - - // 8. Prepare a fixed-width integers collection - ExampleWithFixedWidthIntegersCollection fixedWidthInts; - auto fixedWidthObj = fixedWidthInts.create(); - fixedWidthObj.fixedI16(-12345); - fixedWidthObj.fixedU32(1234567890U); - fixedWidthObj.fixedU64(1234567890123456789ULL); - - auto& maxComp = fixedWidthObj.fixedWidthStruct(); - maxComp.fixedUnsigned16 = 65535; - maxComp.fixedInteger64 = 9223372036854775807LL; - maxComp.fixedInteger32 = 2147483647; - - std::array arrVal = {-10, 20}; - fixedWidthObj.fixedWidthArray(arrVal); - - // Put them all into the frame - originalFrame.put(std::move(hits), "Hits"); - originalFrame.put(std::move(clusters), "Clusters"); - originalFrame.put(std::move(relColls), "OneRelation"); - originalFrame.put(std::move(vecColls), "VectorMember"); - originalFrame.put(std::move(subsetHits), "SubsetHits"); - originalFrame.put(std::move(arrays), "ExampleWithArray"); - originalFrame.put(std::move(emptyClusterColl), "EmptyCluster"); - originalFrame.put(std::move(fixedWidthInts), "FixedWidthInts"); - - // Put some frame parameters - originalFrame.putParameter("anInt", 42); - originalFrame.putParameter("someFloats", std::vector{1.23f, 2.34f, 3.45f}); - originalFrame.putParameter("someStrings", std::vector{"hello", "world"}); - - std::vector colls = {"Hits", "Clusters", "OneRelation", "VectorMember", - "SubsetHits", "ExampleWithArray", "EmptyCluster", "FixedWidthInts"}; +TEST_CASE("ArrowFrameConverter - Comprehensive Round-Trip (No-UserData)", "[arrow][converter][common]") { + auto originalFrame = makeFrame(0); + + const std::vector colls = {"mcparticles", + "moreMCs", + "arrays", + "mcParticleRefs", + "hits", + "hitRefs", + "refs", + "refs2", + "clusters", + "OneRelation", + "info", + "WithVectorMember", + "VectorMemberSubsetColl", + "fixedWidthInts", + "WithNamespaceMember", + "WithNamespaceRelation", + "WithNamespaceRelationCopy", + "emptyCollection", + "emptySubsetColl", + "extension_Contained", + "extension_ExternalComponent", + "extension_ExternalRelation", + "interface_examples", + "anotherHits", + "extension_interface_relation", + "links", + "links_with_interfaces", + "extension_interface_links"}; - // Convert to Arrow Table auto table = podio::convertFrameToTable(originalFrame, colls); REQUIRE(table != nullptr); - // Convert Arrow Table to JSON - nlohmann::json arrowJson = arrowTableToJson(table); - - // Convert Frame to JSON directly - nlohmann::json frameJson = nlohmann::json::object(); - frameJson["Hits"] = originalFrame.get("Hits"); - frameJson["Clusters"] = originalFrame.get("Clusters"); - frameJson["OneRelation"] = originalFrame.get("OneRelation"); - frameJson["VectorMember"] = originalFrame.get("VectorMember"); - frameJson["SubsetHits"] = originalFrame.get("SubsetHits"); - frameJson["ExampleWithArray"] = originalFrame.get("ExampleWithArray"); - frameJson["EmptyCluster"] = originalFrame.get("EmptyCluster"); - frameJson["FixedWidthInts"] = originalFrame.get("FixedWidthInts"); - frameJson["frame_parameters"] = genericParametersToJson(originalFrame.getParameters()); - - // Compare the JSON representations - REQUIRE(arrowJson == frameJson); + auto reconstructedFrame = podio::convertTableToFrame(table, 0); + + verifyEventNoUserData(reconstructedFrame, 0); + + processExtensions(reconstructedFrame, 0, podio::version::build_version); + checkVecMemSubsetColl(reconstructedFrame); + checkInterfaceCollection(reconstructedFrame); + checkInterfaceExtension(reconstructedFrame); + + const auto& hits = reconstructedFrame.get("hits"); + const auto& clusters = reconstructedFrame.get("clusters"); + checkLinkCollection(reconstructedFrame, hits, clusters); + + // Verify Link collection with interfaces + const auto& interfaceLinks = reconstructedFrame.get("links_with_interfaces"); + REQUIRE(interfaceLinks.size() == 3); + const auto& mcps = reconstructedFrame.get("mcparticles"); + REQUIRE(interfaceLinks[0].get() == clusters[0]); + REQUIRE(interfaceLinks[0].get() == hits[0]); + REQUIRE(interfaceLinks[1].get() == clusters[1]); + REQUIRE(interfaceLinks[1].get() == mcps[0]); + REQUIRE(interfaceLinks[2].get() == clusters[0]); + REQUIRE(interfaceLinks[2].get() == clusters[1]); +} + +TEST_CASE("ArrowFrameConverter - convertTableToFrame Verification (Multi-Row / rowIndex > 0)", + "[arrow][converter][reader]") { + auto frame1 = makeFrame(0); + auto frame2 = makeFrame(1); + + const std::vector colls = {"mcparticles", + "moreMCs", + "arrays", + "mcParticleRefs", + "hits", + "hitRefs", + "refs", + "refs2", + "clusters", + "OneRelation", + "info", + "WithVectorMember", + "VectorMemberSubsetColl", + "fixedWidthInts", + "WithNamespaceMember", + "WithNamespaceRelation", + "WithNamespaceRelationCopy", + "emptyCollection", + "emptySubsetColl", + "extension_Contained", + "extension_ExternalComponent", + "extension_ExternalRelation", + "interface_examples", + "anotherHits", + "extension_interface_relation", + "links", + "links_with_interfaces", + "extension_interface_links"}; + + // Convert both frames to Arrow Tables + auto table1 = podio::convertFrameToTable(frame1, colls); + REQUIRE(table1 != nullptr); + auto table2 = podio::convertFrameToTable(frame2, colls); + REQUIRE(table2 != nullptr); + + // Concatenate tables to simulate a multi-row table (e.g. from multiple events) + auto concatResult = arrow::ConcatenateTables({table1, table2}); + REQUIRE(concatResult.ok()); + const auto& concatTable = concatResult.ValueOrDie(); + REQUIRE(concatTable->num_rows() == 2); + + auto verifyFrame = [](const podio::Frame& reconstructedFrame, int eventNum) { + verifyEventNoUserData(reconstructedFrame, eventNum); + + processExtensions(reconstructedFrame, eventNum, podio::version::build_version); + checkVecMemSubsetColl(reconstructedFrame); + checkInterfaceCollection(reconstructedFrame); + checkInterfaceExtension(reconstructedFrame); + + const auto& hits = reconstructedFrame.get("hits"); + const auto& clusters = reconstructedFrame.get("clusters"); + checkLinkCollection(reconstructedFrame, hits, clusters); + + // Verify Link collection with interfaces + const auto& interfaceLinks = reconstructedFrame.get("links_with_interfaces"); + REQUIRE(interfaceLinks.size() == 3); + const auto& mcps = reconstructedFrame.get("mcparticles"); + REQUIRE(interfaceLinks[0].get() == clusters[0]); + REQUIRE(interfaceLinks[0].get() == hits[0]); + REQUIRE(interfaceLinks[1].get() == clusters[1]); + REQUIRE(interfaceLinks[1].get() == mcps[0]); + REQUIRE(interfaceLinks[2].get() == clusters[0]); + REQUIRE(interfaceLinks[2].get() == clusters[1]); + }; + + // --- Reconstruct and verify Frame 1 at rowIndex = 0 --- + auto reconstructedFrame1 = podio::convertTableToFrame(concatTable, 0); + verifyFrame(reconstructedFrame1, 0); + + // --- Reconstruct and verify Frame 2 at rowIndex = 1 --- + auto reconstructedFrame2 = podio::convertTableToFrame(concatTable, 1); + verifyFrame(reconstructedFrame2, 1); }