feat: Add Apache Arrow schema mapper generation to code generator (#966)

arnavdham · web-flow · commit 9bcd402e6e89 · 2026-06-09T19:06:03.000Z
diff --git a/.github/workflows/sanitizers.yaml b/.github/workflows/sanitizers.yaml
@@ -44,6 +44,7 @@ jobs:
             cmake -B build -S . \
               --preset ci-build \
               -DCMAKE_BUILD_TYPE=Debug \
+              -DENABLE_ARROW=ON \
               -DUSE_SANITIZER=${{ matrix.sanitizer }} \
               -DCMAKE_CXX_STANDARD=$([[ ${{ matrix.compiler }} == gcc15 ]] && echo "23" || echo "20")
             echo "::endgroup::"
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -36,6 +36,7 @@ jobs:
           # export JULIA_DEPOT_PATH="$(mktemp -d -p /tmp -t julia_depot_XXXXX):"
           cmake --preset ci-build \
             -DENABLE_JULIA=ON \
+            -DENABLE_ARROW=ON \
             -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/.github/workflows/ubuntu.yml b/.github/workflows/ubuntu.yml
@@ -32,6 +32,7 @@ jobs:
           echo "::group::Run CMake"
           cmake -S . --preset ci-build \
             -DENABLE_JULIA=ON \
+            -DENABLE_ARROW=ON \
             -DCMAKE_INSTALL_PREFIX=../install \
             -DCMAKE_CXX_STANDARD=20 \
             -DUSE_EXTERNAL_CATCH2=OFF \
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -67,6 +67,7 @@ ADD_CLANG_TIDY()
 #--- Declare options -----------------------------------------------------------
 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(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)
@@ -146,6 +147,15 @@ if(ENABLE_SIO)
   endif()
 endif()
 
+# optionally build with Arrow -----------------------------------------------
+if(ENABLE_ARROW)
+  find_package(Arrow REQUIRED)
+  if(Arrow_FOUND)
+    message(STATUS "Found Arrow library - will build Arrow I/O support")
+    list(APPEND PODIO_IO_HANDLERS ARROW)
+  endif()
+endif()
+
 #--- enable unit testing capabilities ------------------------------------------
 include(CTest)
 set(USE_EXTERNAL_CATCH2 AUTO CACHE STRING "Link against an external Catch2 v3 static library, otherwise build it locally")
diff --git a/cmake/podioMacros.cmake b/cmake/podioMacros.cmake
@@ -197,6 +197,7 @@ function(PODIO_ADD_DATAMODEL_CORE_LIB lib_name HEADERS SOURCES)
 
   # Filter out anything I/O backend related to build the core library
   LIST(FILTER HEADERS EXCLUDE REGEX .*SIOBlock.h)
+  LIST(FILTER HEADERS EXCLUDE REGEX .*ArrowMapper.h)
   LIST(FILTER SOURCES EXCLUDE REGEX .*SIOBlock.cc)
 
   add_library(${lib_name} SHARED ${SOURCES} ${HEADERS})
@@ -250,6 +251,7 @@ function(PODIO_ADD_ROOT_IO_DICT dict_name CORE_LIB HEADERS SELECTION_XML)
   # Filter out anything I/O backend related from the generated headers as ROOT only needs
   # the core headers
   LIST(FILTER HEADERS EXCLUDE REGEX .*SIOBlock.h)
+  LIST(FILTER HEADERS EXCLUDE REGEX .*ArrowMapper.h)
 
   add_library(${dict_name} SHARED)
   target_link_libraries(${dict_name} PUBLIC
diff --git a/python/podio_class_generator.py b/python/podio_class_generator.py
@@ -105,7 +105,7 @@ def parse_version(version_str):
     parser.add_argument("packagename", help="Name of the package.")
     parser.add_argument(
         "iohandlers",
-        choices=["ROOT", "SIO"],
+        choices=["ROOT", "SIO", "ARROW"],
         nargs="*",
         help="The IO backend specific code that should be generated",
         default="ROOT",
diff --git a/python/podio_gen/cpp_generator.py b/python/podio_gen/cpp_generator.py
@@ -15,7 +15,7 @@
 )
 from podio_gen.podio_config_reader import PodioConfigReader
 from podio_gen.generator_base import ClassGeneratorBaseMixin, write_file_if_changed
-from podio_gen.generator_utils import DataType, DataModelJSONEncoder
+from podio_gen.generator_utils import DataType, DataModelJSONEncoder, DefinitionError
 
 REPORT_TEXT = """
   PODIO Data Model
@@ -24,6 +24,34 @@
   Read instructions in the README.md to run your first example!
 """
 
+ARROW_PRIMITIVE_TYPES = {
+    "bool": "arrow::boolean()",
+    "char": "arrow::int8()",
+    "short": "arrow::int16()",
+    "int": "arrow::int32()",
+    "long": "arrow::int64()",
+    "long long": "arrow::int64()",
+    "unsigned": "arrow::uint32()",
+    "unsigned int": "arrow::uint32()",
+    "unsigned long": "arrow::uint64()",
+    "unsigned long long": "arrow::uint64()",
+    "float": "arrow::float32()",
+    "double": "arrow::float64()",
+    "int16_t": "arrow::int16()",
+    "int32_t": "arrow::int32()",
+    "int64_t": "arrow::int64()",
+    "uint16_t": "arrow::uint16()",
+    "uint32_t": "arrow::uint32()",
+    "uint64_t": "arrow::uint64()",
+    "std::int16_t": "arrow::int16()",
+    "std::int32_t": "arrow::int32()",
+    "std::int64_t": "arrow::int64()",
+    "std::uint16_t": "arrow::uint16()",
+    "std::uint32_t": "arrow::uint32()",
+    "std::uint64_t": "arrow::uint64()",
+    "std::string": "arrow::utf8()",
+}
+
 
 class IncludeFrom(IntEnum):
     """Enum to signify if an include is needed and from where it should come"""
@@ -85,6 +113,9 @@ def post_process(self, datamodel):
         if "ROOT" in self.io_handlers:
             self._create_selection_xml()
 
+        if "ARROW" in self.io_handlers:
+            self._write_arrow_mapper_header(datamodel)
+
         if the_links := datamodel["links"]:
             self._write_links_registration_file(the_links)
         self._write_all_collections_header()
@@ -122,6 +153,76 @@ def do_process_datatype(self, name, datatype):
 
         return datatype
 
+    def _write_arrow_mapper_header(self, datamodel):
+        """A generated helper that exposes the datamodel as an Arrow schema"""
+        datatypes = []
+        for datatype in datamodel["datatypes"]:
+            datatype["arrow_fields"] = self._arrow_fields(datatype)
+            datatypes.append(datatype)
+
+        data = {
+            "package_name": self.package_name,
+            "schema_version": self.datamodel.schema_version,
+            "datatypes": datatypes,
+        }
+        self._write_file(
+            "ArrowMapper.h",
+            self._eval_template("ArrowMapper.h.jinja2", data),
+        )
+
+    def _arrow_fields(self, datatype):
+        """Create Arrow field expressions for the members and relations of a datatype"""
+        fields = []
+        fields.extend(
+            self._arrow_field(member.name, self._arrow_type(member))
+            for member in datatype["Members"]
+        )
+        fields.extend(
+            self._arrow_field(member.name, f"arrow::list({self._arrow_type(member)})")
+            for member in datatype["VectorMembers"]
+        )
+        fields.extend(
+            self._arrow_field(relation.name, "objectRefType()")
+            for relation in datatype["OneToOneRelations"]
+        )
+        fields.extend(
+            self._arrow_field(relation.name, "arrow::list(objectRefType())")
+            for relation in datatype["OneToManyRelations"]
+        )
+        return fields
+
+    def _arrow_field(self, name, type_expr, nullable=False):
+        """Create a C++ arrow::field expression"""
+        nullable_arg = ", true" if nullable else ""
+        return f'arrow::field("{name}", {type_expr}{nullable_arg})'
+
+    def _arrow_type(self, member):
+        """Map a parsed podio member to an Arrow C++ DataType expression"""
+        if member.is_array:
+            value_type = self._arrow_type_from_name(member.array_type)
+            return f"arrow::fixed_size_list({value_type}, {member.array_size})"
+
+        return self._arrow_type_from_name(member.full_type)
+
+    def _arrow_type_from_name(self, type_name):
+        """Map a C++ type name from the datamodel to an Arrow C++ DataType expression"""
+        type_name = type_name.removeprefix("::")
+        if type_name in ARROW_PRIMITIVE_TYPES:
+            return ARROW_PRIMITIVE_TYPES[type_name]
+
+        if type_name in self.datamodel.components:
+            return self._arrow_struct_type(self.datamodel.components[type_name]["Members"])
+
+        if self.upstream_edm and type_name in self.upstream_edm.components:
+            return self._arrow_struct_type(self.upstream_edm.components[type_name]["Members"])
+
+        raise DefinitionError(f"Cannot map '{type_name}' to an Arrow type")
+
+    def _arrow_struct_type(self, members):
+        """Create an Arrow struct expression for a component definition"""
+        fields = [self._arrow_field(member.name, self._arrow_type(member)) for member in members]
+        return "arrow::struct_({" + ", ".join(fields) + "})"
+
     def do_process_interface(self, _, interface):
         """Process an interface definition and generate the necessary code"""
         interface["include_types"] = [
diff --git a/python/templates/ArrowMapper.h.jinja2 b/python/templates/ArrowMapper.h.jinja2
@@ -0,0 +1,49 @@
+// AUTOMATICALLY GENERATED FILE - DO NOT EDIT
+
+#ifndef {{ package_name.upper() }}_ARROWMAPPER_H
+#define {{ package_name.upper() }}_ARROWMAPPER_H
+
+#include <arrow/type.h>
+#include <arrow/util/key_value_metadata.h>
+
+#include <memory>
+
+namespace {{ package_name }}::arrow_io {
+
+inline std::shared_ptr<arrow::DataType> objectRefType() {
+  return arrow::struct_({
+      arrow::field("collectionID", arrow::uint32(), false),
+      arrow::field("index", arrow::int32(), false),
+  });
+}
+
+inline std::shared_ptr<arrow::DataType> frameParametersType() {
+  return arrow::struct_({
+      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()))),
+  });
+}
+
+inline std::shared_ptr<arrow::Schema> schema() {
+  return arrow::schema({
+{% for datatype in datatypes %}
+      arrow::field("{{ datatype.class.bare_type }}", arrow::list(arrow::struct_({
+{% for field in datatype.arrow_fields %}
+          {{ field }},
+{% endfor %}
+      }))),
+{% endfor %}
+      arrow::field("frame_parameters", frameParametersType()),
+  });
+}
+
+inline std::shared_ptr<arrow::Schema>
+schemaWithMetadata(std::shared_ptr<const arrow::KeyValueMetadata> metadata) {
+  return schema()->WithMetadata(std::move(metadata));
+}
+
+} // namespace {{ package_name }}::arrow_io
+
+#endif
diff --git a/python/templates/CMakeLists.txt b/python/templates/CMakeLists.txt
@@ -14,6 +14,7 @@ set(PODIO_TEMPLATES
   ${CMAKE_CURRENT_LIST_DIR}/Interface.h.jinja2
   ${CMAKE_CURRENT_LIST_DIR}/MutableObject.cc.jinja2
   ${CMAKE_CURRENT_LIST_DIR}/MutableObject.h.jinja2
+  ${CMAKE_CURRENT_LIST_DIR}/ArrowMapper.h.jinja2
   ${CMAKE_CURRENT_LIST_DIR}/selection.xml.jinja2
   ${CMAKE_CURRENT_LIST_DIR}/SIOBlock.cc.jinja2
   ${CMAKE_CURRENT_LIST_DIR}/SIOBlock.h.jinja2
diff --git a/tests/unittests/CMakeLists.txt b/tests/unittests/CMakeLists.txt
@@ -50,6 +50,10 @@ target_compile_options(unittest_podio PRIVATE -Wno-c2y-extensions)
 if (ENABLE_SIO)
   target_link_libraries(unittest_podio PRIVATE podio::podioSioIO)
 endif()
+if (ENABLE_ARROW)
+  target_link_libraries(unittest_podio PRIVATE Arrow::arrow_shared)
+  target_compile_definitions(unittest_podio PRIVATE PODIO_ENABLE_ARROW=1)
+endif()
 
 # The unittests can easily be filtered and they are labelled so we can put together a
 # list of labels that we want to ignore
diff --git a/tests/unittests/unittest.cpp b/tests/unittests/unittest.cpp
@@ -1778,3 +1778,103 @@ TEST_CASE("Add type lists", "[basics][code-gen]") {
                    "podio::UserDataCollection<uint8_t>", "podio::UserDataCollection<uint16_t>",
                    "podio::UserDataCollection<uint32_t>", "podio::UserDataCollection<uint64_t>"}));
 }
+
+#if PODIO_ENABLE_ARROW
+  #include "datamodel/ArrowMapper.h"
+  #include <arrow/type.h>
+
+TEST_CASE("Arrow schema mapping", "[arrow]") {
+  auto schema = datamodel::arrow_io::schema();
+  REQUIRE(schema != nullptr);
+
+  // 1. Verify top-level fields
+  auto field_names = schema->field_names();
+  REQUIRE(std::find(field_names.begin(), field_names.end(), "EventInfo") != field_names.end());
+  REQUIRE(std::find(field_names.begin(), field_names.end(), "ExampleHit") != field_names.end());
+  REQUIRE(std::find(field_names.begin(), field_names.end(), "ExampleMC") != field_names.end());
+  REQUIRE(std::find(field_names.begin(), field_names.end(), "ExampleWithArray") != field_names.end());
+  REQUIRE(std::find(field_names.begin(), field_names.end(), "frame_parameters") != field_names.end());
+
+  // 2. Verify EventInfo (Simple type)
+  auto eventInfoField = schema->GetFieldByName("EventInfo");
+  REQUIRE(eventInfoField != nullptr);
+  REQUIRE(eventInfoField->type()->id() == arrow::Type::LIST);
+  auto eventInfoList = std::static_pointer_cast<arrow::ListType>(eventInfoField->type());
+  auto eventInfoStruct = std::static_pointer_cast<arrow::StructType>(eventInfoList->value_type());
+  REQUIRE(eventInfoStruct != nullptr);
+  auto numberField = eventInfoStruct->GetFieldByName("Number");
+  REQUIRE(numberField != nullptr);
+  REQUIRE(numberField->type()->id() == arrow::Type::INT32);
+
+  // 3. Verify ExampleHit (Primitive and struct mapping)
+  auto exampleHitField = schema->GetFieldByName("ExampleHit");
+  REQUIRE(exampleHitField != nullptr);
+  REQUIRE(exampleHitField->type()->id() == arrow::Type::LIST);
+  auto exampleHitList = std::static_pointer_cast<arrow::ListType>(exampleHitField->type());
+  auto exampleHitStruct = std::static_pointer_cast<arrow::StructType>(exampleHitList->value_type());
+  REQUIRE(exampleHitStruct != nullptr);
+
+  auto cellIDField = exampleHitStruct->GetFieldByName("cellID");
+  REQUIRE(cellIDField != nullptr);
+  REQUIRE(cellIDField->type()->id() == arrow::Type::UINT64);
+
+  auto energyField = exampleHitStruct->GetFieldByName("energy");
+  REQUIRE(energyField != nullptr);
+  REQUIRE(energyField->type()->id() == arrow::Type::DOUBLE);
+
+  auto xField = exampleHitStruct->GetFieldByName("x");
+  REQUIRE(xField != nullptr);
+  REQUIRE(xField->type()->id() == arrow::Type::DOUBLE);
+
+  // 4. Verify ExampleWithArray (Fixed size lists)
+  auto exampleWithArrayField = schema->GetFieldByName("ExampleWithArray");
+  REQUIRE(exampleWithArrayField != nullptr);
+  auto exampleWithArrayList = std::static_pointer_cast<arrow::ListType>(exampleWithArrayField->type());
+  auto exampleWithArrayStruct = std::static_pointer_cast<arrow::StructType>(exampleWithArrayList->value_type());
+  REQUIRE(exampleWithArrayStruct != nullptr);
+
+  auto myArrayField = exampleWithArrayStruct->GetFieldByName("myArray");
+  REQUIRE(myArrayField != nullptr);
+  REQUIRE(myArrayField->type()->id() == arrow::Type::FIXED_SIZE_LIST);
+  auto myArrayType = std::static_pointer_cast<arrow::FixedSizeListType>(myArrayField->type());
+  REQUIRE(myArrayType->list_size() == 4);
+  REQUIRE(myArrayType->value_type()->id() == arrow::Type::INT32);
+
+  // 5. Verify ExampleMC (Relations mapped to object references)
+  auto exampleMCField = schema->GetFieldByName("ExampleMC");
+  REQUIRE(exampleMCField != nullptr);
+  auto exampleMCList = std::static_pointer_cast<arrow::ListType>(exampleMCField->type());
+  auto exampleMCStruct = std::static_pointer_cast<arrow::StructType>(exampleMCList->value_type());
+  REQUIRE(exampleMCStruct != nullptr);
+
+  auto parentsField = exampleMCStruct->GetFieldByName("parents");
+  REQUIRE(parentsField != nullptr);
+  REQUIRE(parentsField->type()->id() == arrow::Type::LIST);
+  auto parentsList = std::static_pointer_cast<arrow::ListType>(parentsField->type());
+  auto parentsStruct = std::static_pointer_cast<arrow::StructType>(parentsList->value_type());
+  REQUIRE(parentsStruct != nullptr);
+
+  auto parentCollId = parentsStruct->GetFieldByName("collectionID");
+  REQUIRE(parentCollId != nullptr);
+  REQUIRE(parentCollId->type()->id() == arrow::Type::UINT32);
+
+  auto parentIndex = parentsStruct->GetFieldByName("index");
+  REQUIRE(parentIndex != nullptr);
+  REQUIRE(parentIndex->type()->id() == arrow::Type::INT32);
+
+  // 6. Verify frame_parameters (Map mapping)
+  auto frameParamsField = schema->GetFieldByName("frame_parameters");
+  REQUIRE(frameParamsField != nullptr);
+  REQUIRE(frameParamsField->type()->id() == arrow::Type::STRUCT);
+
+  auto frameParamsType = std::static_pointer_cast<arrow::StructType>(frameParamsField->type());
+  auto intParamsField = frameParamsType->GetFieldByName("int_params");
+  REQUIRE(intParamsField != nullptr);
+  REQUIRE(intParamsField->type()->id() == arrow::Type::MAP);
+  auto intParamsType = std::static_pointer_cast<arrow::MapType>(intParamsField->type());
+  REQUIRE(intParamsType->key_type()->id() == arrow::Type::STRING);
+  REQUIRE(intParamsType->item_type()->id() == arrow::Type::LIST);
+  auto intParamsList = std::static_pointer_cast<arrow::ListType>(intParamsType->item_type());
+  REQUIRE(intParamsList->value_type()->id() == arrow::Type::INT32);
+}
+#endif
