index mapping serialization

Author: aalkin

Framework/AnalysisSupport/src/AODReaderHelpers.cxx

@@ -18,20 +18,22 @@
 #include "Framework/CallbackService.h"
 #include "Framework/DataSpecUtils.h"
 #include "../src/ExpressionJSONHelpers.h"
+#include "../src/IndexJSONHelpers.h"
 #include "Framework/ConfigContext.h"
 #include "Framework/AnalysisContext.h"
 
 namespace o2::framework::readers
 {
+namespace
+{
 template <size_t N, std::array<soa::TableRef, N> refs>
 static inline auto extractOriginals(ProcessingContext& pc)
 {
   return [&]<size_t... Is>(std::index_sequence<Is...>) -> std::vector<std::shared_ptr<arrow::Table>> {
     return {pc.inputs().get<TableConsumer>(o2::aod::label<refs[Is]>())->asArrowTable()...};
   }(std::make_index_sequence<refs.size()>());
 }
-namespace
-{
+
 template <typename D>
   requires(D::exclusive)
 auto make_build(D metadata, InputSpec const& input, ProcessingContext& pc)
@@ -58,16 +60,39 @@ auto make_build(D metadata, InputSpec const& input, ProcessingContext& pc)
                                                                                                         index_pack_t{});
 }
 
+static inline auto extractSources(ProcessingContext& pc, std::vector<std::string> const& labels)
+{
+  std::vector<std::shared_ptr<arrow::Table>> tables;
+  for (auto const& label : labels) {
+    tables.emplace_back(pc.inputs().get<TableConsumer>(label.c_str())->asArrowTable());
+  }
+  return tables;
+}
+
 struct Builder {
+  std::string binding;
+  std::vector<std::string> labels;
+  std::vector<o2::soa::IndexRecord> records;
+  header::DataOrigin origin;
+  header::DataDescription description;
+  header::DataHeader::SubSpecificationType version;
+
+  std::shared_ptr<arrow::Table> build(ProcessingContext& pc) const
+  {
+    std::shared_ptr<arrow::Table> result;
+    auto tables = extractSources(pc, labels);
+    return result;
+  }
 
 };
 
 struct Buildable {
   std::string binding;
-
+  std::vector<std::string> labels;
   header::DataOrigin origin;
   header::DataDescription description;
   header::DataHeader::SubSpecificationType version;
+  std::vector<o2::soa::IndexRecord> records;
 
   Buildable(InputSpec const& spec)
     : binding{spec.binding}
@@ -77,14 +102,25 @@ struct Buildable {
     description = description_;
     version = version_;
 
-    // The following components are needed to build an index table
-    // 1. the labels of the source tables to extract from inputRecord -> extracted from input metadata
-    // 2. the mapping, in the order of the definition of columns, of the
-    //    position in each source table of an index column pointing to the Key
-    //    and the types of index to write (self, single-valued, slice or array)
-    // the mapping has to be created at the point where the type information is available and
-    // put into the input spec metadata as a vector of (type, label, pos)
+    auto loc = std::find_if(spec.metadata.begin(), spec.metadata.end(), [](ConfigParamSpec const& cps) { return cps.name.compare("index-records") == 0; });
+    std::stringstream iws(loc->defaultValue.get<std::string>());
+    records = IndexJSONHelpers::read(iws);
 
+    for (auto const& r : records) {
+      labels.emplace_back(r.label);
+    }
+  }
+
+  Builder createBuilder() const
+  {
+    return Builder{
+      binding,
+      labels,
+      records,
+      origin,
+      description,
+      version
+    };
   }
 
 };
@@ -132,21 +168,6 @@ AlgorithmSpec AODReaderHelpers::indexBuilderCallback(ConfigContext const& ctx)
 
 namespace
 {
-template <o2::aod::is_aod_hash D>
-auto make_spawn(InputSpec const& input, ProcessingContext& pc)
-{
-  using metadata_t = o2::aod::MetadataTrait<D>::metadata;
-  constexpr auto sources = metadata_t::sources;
-  static std::shared_ptr<gandiva::Projector> projector = nullptr;
-  static std::shared_ptr<arrow::Schema> schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(typename metadata_t::expression_pack_t{}));
-  static auto projectors = []<typename... C>(framework::pack<C...>) -> std::array<expressions::Projector, sizeof...(C)>
-  {
-    return {{std::move(C::Projector())...}};
-  }
-  (typename metadata_t::expression_pack_t{});
-  return o2::framework::spawner<D>(extractOriginals<sources.size(), sources>(pc), input.binding.c_str(), projectors.data(), projector, schema);
-}
-
 struct Maker {
   std::string binding;
   std::vector<std::string> labels;

Framework/Core/CMakeLists.txt

@@ -142,6 +142,7 @@ o2_add_library(Framework
                        src/Variant.cxx
                        src/VariantJSONHelpers.cxx
                        src/ExpressionJSONHelpers.cxx
+		       src/IndexJSONHelpers.cxx
                        src/VariantPropertyTreeHelpers.cxx
                        src/WorkflowCustomizationHelpers.cxx
                        src/WorkflowHelpers.cxx

Framework/Core/include/Framework/ASoA.h

@@ -217,6 +217,19 @@ using is_self_index_t = typename std::conditional_t<is_self_index_column<C>, std
 
 namespace o2::aod
 {
+namespace {
+template <typename Key, size_t N, std::array<bool, N> map>
+static consteval int getIndexPosToKey_impl()
+{
+  constexpr const auto pos = std::find(map.begin(), map.end(), true);
+  if constexpr (pos != map.end()) {
+    return std::distance(map.begin(), pos);
+  } else {
+    return -1;
+  }
+}
+}
+
 /// Base type for table metadata
 template <typename D, typename... Cs>
 struct TableMetadata {
@@ -243,17 +256,6 @@ struct TableMetadata {
     return getIndexPosToKey_impl<Key, framework::pack_size(persistent_columns_t{}), getMap<Key>(persistent_columns_t{})>();
   }
 
-  template <typename Key, size_t N, std::array<bool, N> map>
-  static consteval int getIndexPosToKey_impl()
-  {
-    constexpr const auto pos = std::find(map.begin(), map.end(), true);
-    if constexpr (pos != map.end()) {
-      return std::distance(map.begin(), pos);
-    } else {
-      return -1;
-    }
-  }
-
   static std::shared_ptr<arrow::Schema> getSchema()
   {
     return std::make_shared<arrow::Schema>([]<typename... C>(framework::pack<C...>&& p) { return o2::soa::createFieldsFromColumns(p); }(persistent_columns_t{}));

Framework/Core/include/Framework/AnalysisHelpers.h

@@ -25,11 +25,30 @@
 #include "Framework/TableBuilder.h"
 #include "Framework/Traits.h"
 
+#include <cstdio>
 #include <string>
+namespace o2::soa {
+enum struct IndexKind : int {
+  IdxInvalid = -1,
+  IdxSelf = 0,
+  IdxSingle = 1,
+  IdxSlice = 2,
+  IdxArray = 3
+};
+
+struct IndexRecord {
+  std::string label;
+  IndexKind kind;
+  int pos;
+  auto operator<=>(const IndexRecord&) const = default;
+};
+} // namespace o2::soa
+
 namespace o2::framework
 {
 std::string serializeProjectors(std::vector<framework::expressions::Projector>& projectors);
 std::string serializeSchema(std::shared_ptr<arrow::Schema> schema);
+std::string serializeIndexRecords(std::vector<o2::soa::IndexRecord>& irs);
 }  // namespace o2::framework
 
 namespace o2::soa
@@ -54,19 +73,6 @@ constexpr auto tableRef2Schema()
     {"\"\""}};
 }
 
-enum struct IndexKind : unsigned {
-  IdxSelf = 0,
-  IdxSingle = 1,
-  IdxSlice = 2,
-  IdxArray = 3
-};
-
-struct IndexRecord {
-  std::string label;
-  IndexKind kind;
-  size_t pos;
-};
-
 namespace
 {
 template <soa::with_sources T>
@@ -102,6 +108,47 @@ inline constexpr auto getCCDBUrls()
   return result;
 }
 
+template <typename T>
+  requires(std::same_as<T, int>)
+consteval IndexKind getIndexKind()
+{
+  return IndexKind::IdxSingle;
+}
+
+template <typename T>
+  requires(std::is_bounded_array_v<T>)
+consteval IndexKind getIndexKind()
+{
+  return IndexKind::IdxSlice;
+}
+
+template <typename T>
+  requires(framework::is_specialization_v<T, std::vector>)
+consteval IndexKind getIndexKind()
+{
+  return IndexKind::IdxArray;
+}
+
+template <soa::with_index_pack T>
+inline constexpr auto getIndexMapping()
+{
+  std::vector<IndexRecord> idx;
+  using indices = T::index_pack_t;
+  using Key = T::Key;
+  [&idx]<size_t... Is>(std::index_sequence<Is...>) mutable {
+    constexpr auto refs = T::sources;
+    ([&idx]<TableRef ref, typename CT>() mutable {
+      constexpr auto pos = o2::aod::MetadataTrait<o2::aod::Hash<ref.desc_hash>>::metadata::template getIndexPosToKey<Key>();
+      if constexpr (pos == -1) {
+        idx.emplace_back(o2::aod::label<ref>(), IndexKind::IdxSelf, pos);
+      } else {
+        idx.emplace_back(o2::aod::label<ref>(), getIndexKind<CT>(), pos);
+      }
+    }.template operator()<refs[Is], typename framework::pack_element_t<Is, indices>::type>(), ...);
+  }(std::make_index_sequence<framework::pack_size(indices{})>());;
+  return idx;
+}
+
 template <soa::with_sources T>
 constexpr auto getInputMetadata() -> std::vector<framework::ConfigParamSpec>
 {
@@ -170,8 +217,8 @@ constexpr auto getExpressionMetadata() -> std::vector<framework::ConfigParamSpec
 template <soa::with_index_pack T>
 constexpr auto getIndexMetadata() -> std::vector<framework::ConfigParamSpec>
 {
-
-  return {};
+  auto map = getIndexMapping<T>();
+  return {framework::ConfigParamSpec{"index-records", framework::VariantType::String, framework::serializeIndexRecords(map), {"\"\""}}};
 }
 
 template <typename T>
@@ -193,6 +240,8 @@ constexpr auto tableRef2InputSpec()
   metadata.insert(metadata.end(), ccdbMetadata.begin(), ccdbMetadata.end());
   auto p = getExpressionMetadata<typename o2::aod::MetadataTrait<o2::aod::Hash<R.desc_hash>>::metadata>();
   metadata.insert(metadata.end(), p.begin(), p.end());
+  auto idx = getIndexMetadata<typename o2::aod::MetadataTrait<o2::aod::Hash<R.desc_hash>>::metadata>();
+  metadata.insert(metadata.end(), idx.begin(), idx.end());
   if constexpr (!soa::with_ccdb_urls<typename o2::aod::MetadataTrait<o2::aod::Hash<R.desc_hash>>::metadata>) {
     metadata.emplace_back(framework::ConfigParamSpec{"schema", framework::VariantType::String, framework::serializeSchema(o2::aod::MetadataTrait<o2::aod::Hash<R.desc_hash>>::metadata::getSchema()), {"\"\""}});
   }

Framework/Core/src/AnalysisHelpers.cxx

@@ -8,8 +8,10 @@
 // In applying this license CERN does not waive the privileges and immunities
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
+#include "Framework/AnalysisHelpers.h"
 #include "Framework/ExpressionHelpers.h"
 #include "ExpressionJSONHelpers.h"
+#include "IndexJSONHelpers.h"
 
 namespace o2::framework
 {
@@ -41,4 +43,11 @@ std::string serializeSchema(std::shared_ptr<arrow::Schema> schema)
   ArrowJSONHelpers::write(osm, schema);
   return osm.str();
 }
+
+std::string serializeIndexRecords(std::vector<o2::soa::IndexRecord>& irs)
+{
+  std::stringstream osm;
+  IndexJSONHelpers::write(osm, irs);
+  return osm.str();
+}
 } // namespace o2::framework