From 9f4fe3bba767245721d4fac46f5cab9b7c4a5379 Mon Sep 17 00:00:00 2001
From: Anton Alkin <anton.alkin@cern.ch>
Date: Thu, 2 Dec 2021 17:29:44 +0100
Subject: [PATCH] DPL Analysis: make sure all supported types can be used in
 vector columns

---
 Framework/Core/include/Framework/ArrowTypes.h | 35 +++++++++++--------
 Framework/Core/test/test_TreeToTable.cxx      | 10 ++++--
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/Framework/Core/include/Framework/ArrowTypes.h b/Framework/Core/include/Framework/ArrowTypes.h
index 474b2611a5eea..7e261b9989aee 100644
--- a/Framework/Core/include/Framework/ArrowTypes.h
+++ b/Framework/Core/include/Framework/ArrowTypes.h
@@ -77,21 +77,26 @@ struct arrow_array_for<double[N]> {
   using type = arrow::FixedSizeListArray;
   using value_type = double;
 };
-template <>
-struct arrow_array_for<std::vector<int>> {
-  using type = arrow::ListArray;
-  using value_type = int;
-};
-template <>
-struct arrow_array_for<std::vector<float>> {
-  using type = arrow::ListArray;
-  using value_type = float;
-};
-template <>
-struct arrow_array_for<std::vector<double>> {
-  using type = arrow::ListArray;
-  using value_type = double;
-};
+
+#define ARROW_VECTOR_FOR(_type_)                \
+  template <>                                   \
+  struct arrow_array_for<std::vector<_type_>> { \
+    using type = arrow::ListArray;              \
+    using value_type = _type_;                  \
+  };
+
+ARROW_VECTOR_FOR(uint8_t);
+ARROW_VECTOR_FOR(uint16_t);
+ARROW_VECTOR_FOR(uint32_t);
+ARROW_VECTOR_FOR(uint64_t);
+
+ARROW_VECTOR_FOR(int8_t);
+ARROW_VECTOR_FOR(int16_t);
+ARROW_VECTOR_FOR(int32_t);
+ARROW_VECTOR_FOR(int64_t);
+
+ARROW_VECTOR_FOR(float);
+ARROW_VECTOR_FOR(double);
 
 template <typename T>
 using arrow_array_for_t = typename arrow_array_for<T>::type;
diff --git a/Framework/Core/test/test_TreeToTable.cxx b/Framework/Core/test/test_TreeToTable.cxx
index aed27db0cda57..cd6833d2217cb 100644
--- a/Framework/Core/test/test_TreeToTable.cxx
+++ b/Framework/Core/test/test_TreeToTable.cxx
@@ -166,9 +166,10 @@ namespace cols
 DECLARE_SOA_COLUMN(Ivec, ivec, std::vector<int>);
 DECLARE_SOA_COLUMN(Fvec, fvec, std::vector<float>);
 DECLARE_SOA_COLUMN(Dvec, dvec, std::vector<double>);
+DECLARE_SOA_COLUMN(UIvec, uivec, std::vector<uint8_t>);
 } // namespace cols
 
-DECLARE_SOA_TABLE(Vectors, "AOD", "VECS", o2::soa::Index<>, cols::Ivec, cols::Fvec, cols::Dvec);
+DECLARE_SOA_TABLE(Vectors, "AOD", "VECS", o2::soa::Index<>, cols::Ivec, cols::Fvec, cols::Dvec, cols::UIvec);
 } // namespace o2::aod
 
 BOOST_AUTO_TEST_CASE(VariableLists)
@@ -178,16 +179,19 @@ BOOST_AUTO_TEST_CASE(VariableLists)
   std::vector<int> iv;
   std::vector<float> fv;
   std::vector<double> dv;
+  std::vector<uint8_t> ui;
   for (auto i = 1; i < 11; ++i) {
     iv.clear();
     fv.clear();
     dv.clear();
+    ui.clear();
     for (auto j = 0; j < i; ++j) {
       iv.push_back(j + 2);
       fv.push_back((j + 2) * 0.2134f);
       dv.push_back((j + 4) * 0.192873819237);
+      ui.push_back(j);
     }
-    writer(0, iv, fv, dv);
+    writer(0, iv, fv, dv, ui);
   }
   auto table = b.finalize();
 
@@ -209,10 +213,12 @@ BOOST_AUTO_TEST_CASE(VariableLists)
     auto iv = row.ivec();
     auto fv = row.fvec();
     auto dv = row.dvec();
+    auto uv = row.uivec();
     for (auto j = 0; j < i; ++j) {
       BOOST_CHECK_EQUAL(iv[j], j + 2);
       BOOST_CHECK_EQUAL(fv[j], (j + 2) * 0.2134f);
       BOOST_CHECK_EQUAL(dv[j], (j + 4) * 0.192873819237);
+      BOOST_CHECK_EQUAL(uv[j], j);
     }
     ++i;
   }