DPL Analysis: prevent slice cache from updating when not required by enabled process functions (#14057)

Author: aalkin

Framework/Core/include/Framework/ASoA.h

@@ -1400,10 +1400,10 @@ namespace o2::framework
 
 struct PreslicePolicyBase {
   const std::string binding;
-  StringPair bindingKey;
+  Entry bindingKey;
 
   bool isMissing() const;
-  StringPair const& getBindingKey() const;
+  Entry const& getBindingKey() const;
 };
 
 struct PreslicePolicySorted : public PreslicePolicyBase {
@@ -1428,7 +1428,7 @@ struct PresliceBase : public Policy {
   const std::string binding;
 
   PresliceBase(expressions::BindingNode index_)
-    : Policy{PreslicePolicyBase{{o2::soa::getLabelFromTypeForKey<T, OPT>(std::string{index_.name})}, std::make_pair(o2::soa::getLabelFromTypeForKey<T, OPT>(std::string{index_.name}), std::string{index_.name})}, {}}
+    : Policy{PreslicePolicyBase{{o2::soa::getLabelFromTypeForKey<T, OPT>(std::string{index_.name})}, Entry(o2::soa::getLabelFromTypeForKey<T, OPT>(std::string{index_.name}), std::string{index_.name})}, {}}
   {
   }
 
@@ -1508,7 +1508,7 @@ auto doSliceBy(T const* table, o2::framework::PresliceBase<C, Policy, OPT> const
 {
   if constexpr (OPT) {
     if (container.isMissing()) {
-      missingOptionalPreslice(getLabelFromType<std::decay_t<T>>().data(), container.bindingKey.second.c_str());
+      missingOptionalPreslice(getLabelFromType<std::decay_t<T>>().data(), container.bindingKey.key.c_str());
     }
   }
   uint64_t offset = 0;
@@ -1545,7 +1545,7 @@ auto doSliceBy(T const* table, o2::framework::PresliceBase<C, Policy, OPT> const
 {
   if constexpr (OPT) {
     if (container.isMissing()) {
-      missingOptionalPreslice(getLabelFromType<std::decay_t<T>>().data(), container.bindingKey.second.c_str());
+      missingOptionalPreslice(getLabelFromType<std::decay_t<T>>().data(), container.bindingKey.key.c_str());
     }
   }
   auto selection = container.getSliceFor(value);
@@ -1574,7 +1574,7 @@ auto doFilteredSliceBy(T const* table, o2::framework::PresliceBase<C, framework:
 {
   if constexpr (OPT) {
     if (container.isMissing()) {
-      missingOptionalPreslice(getLabelFromType<T>().data(), container.bindingKey.second.c_str());
+      missingOptionalPreslice(getLabelFromType<T>().data(), container.bindingKey.key.c_str());
     }
   }
   uint64_t offset = 0;

Framework/Core/include/Framework/AnalysisManagers.h

@@ -534,39 +534,43 @@ static void setGroupedCombination(C& comb, TG& grouping, std::tuple<Ts...>& asso
 /// Preslice handling
 template <typename T>
   requires(!is_preslice<T>)
-bool registerCache(T&, std::vector<StringPair>&, std::vector<StringPair>&)
+bool registerCache(T&, Cache&, Cache&)
 {
   return false;
 }
 
 template <is_preslice T>
   requires std::same_as<typename T::policy_t, framework::PreslicePolicySorted>
-bool registerCache(T& preslice, std::vector<StringPair>& bsks, std::vector<StringPair>&)
+bool registerCache(T& preslice, Cache& bsks, Cache&)
 {
   if constexpr (T::optional) {
     if (preslice.binding == "[MISSING]") {
       return true;
     }
   }
-  auto locate = std::find_if(bsks.begin(), bsks.end(), [&](auto const& entry) { return (entry.first == preslice.bindingKey.first) && (entry.second == preslice.bindingKey.second); });
+  auto locate = std::find_if(bsks.begin(), bsks.end(), [&](auto const& entry) { return (entry.binding == preslice.bindingKey.binding) && (entry.key == preslice.bindingKey.key); });
   if (locate == bsks.end()) {
     bsks.emplace_back(preslice.getBindingKey());
+  } else if (locate->enabled == false) {
+    locate->enabled = true;
   }
   return true;
 }
 
 template <is_preslice T>
   requires std::same_as<typename T::policy_t, framework::PreslicePolicyGeneral>
-bool registerCache(T& preslice, std::vector<StringPair>&, std::vector<StringPair>& bsksU)
+bool registerCache(T& preslice, Cache&, Cache& bsksU)
 {
   if constexpr (T::optional) {
     if (preslice.binding == "[MISSING]") {
       return true;
     }
   }
-  auto locate = std::find_if(bsksU.begin(), bsksU.end(), [&](auto const& entry) { return (entry.first == preslice.bindingKey.first) && (entry.second == preslice.bindingKey.second); });
+  auto locate = std::find_if(bsksU.begin(), bsksU.end(), [&](auto const& entry) { return (entry.binding == preslice.bindingKey.binding) && (entry.key == preslice.bindingKey.key); });
   if (locate == bsksU.end()) {
     bsksU.emplace_back(preslice.getBindingKey());
+  } else if (locate->enabled == false) {
+    locate->enabled = true;
   }
   return true;
 }

Framework/Core/include/Framework/AnalysisTask.h

@@ -66,20 +66,20 @@ concept is_enumeration = is_enumeration_v<std::decay_t<T>>;
 namespace {
 struct AnalysisDataProcessorBuilder {
   template <typename G, typename... Args>
-  static void addGroupingCandidates(std::vector<StringPair>& bk, std::vector<StringPair>& bku)
+  static void addGroupingCandidates(Cache& bk, Cache& bku, bool enabled)
   {
-    [&bk, &bku]<typename... As>(framework::pack<As...>) mutable {
+    [&bk, &bku, enabled]<typename... As>(framework::pack<As...>) mutable {
       std::string key;
       if constexpr (soa::is_iterator<std::decay_t<G>>) {
         key = std::string{"fIndex"} + o2::framework::cutString(soa::getLabelFromType<std::decay_t<G>>());
       }
-      ([&bk, &bku, &key]() mutable {
+      ([&bk, &bku, &key, enabled]() mutable {
         if constexpr (soa::relatedByIndex<std::decay_t<G>, std::decay_t<As>>()) {
           auto binding = soa::getLabelFromTypeForKey<std::decay_t<As>>(key);
           if constexpr (o2::soa::is_smallgroups<std::decay_t<As>>) {
-            framework::updatePairList(bku, binding, key);
+            framework::updatePairList(bku, binding, key, enabled);
           } else {
-            framework::updatePairList(bk, binding, key);
+            framework::updatePairList(bk, binding, key, enabled);
           }
         }
       }(),
@@ -147,7 +147,7 @@ struct AnalysisDataProcessorBuilder {
   /// helper to parse the process arguments
   /// 1. enumeration (must be the only argument)
   template <typename R, typename C, is_enumeration A>
-  static void inputsFromArgs(R (C::*)(A), const char* /*name*/, bool /*value*/, std::vector<InputSpec>& inputs, std::vector<ExpressionInfo>&, std::vector<StringPair>&, std::vector<StringPair>&)
+  static void inputsFromArgs(R (C::*)(A), const char* /*name*/, bool /*value*/, std::vector<InputSpec>& inputs, std::vector<ExpressionInfo>&, Cache&, Cache&)
   {
     std::vector<ConfigParamSpec> inputMetadata;
     // FIXME: for the moment we do not support begin, end and step.
@@ -156,17 +156,17 @@ struct AnalysisDataProcessorBuilder {
 
   /// 2. grouping case - 1st argument is an iterator
   template <typename R, typename C, soa::is_iterator A, soa::is_table... Args>
-  static void inputsFromArgs(R (C::*)(A, Args...), const char* name, bool value, std::vector<InputSpec>& inputs, std::vector<ExpressionInfo>& eInfos, std::vector<StringPair>& bk, std::vector<StringPair>& bku)
+  static void inputsFromArgs(R (C::*)(A, Args...), const char* name, bool value, std::vector<InputSpec>& inputs, std::vector<ExpressionInfo>& eInfos, Cache& bk, Cache& bku)
     requires(std::is_lvalue_reference_v<A> && (std::is_lvalue_reference_v<Args> && ...))
   {
-    addGroupingCandidates<A, Args...>(bk, bku);
+    addGroupingCandidates<A, Args...>(bk, bku, value);
     constexpr auto hash = o2::framework::TypeIdHelpers::uniqueId<R (C::*)(A, Args...)>();
     addInputsAndExpressions<typename std::decay_t<A>::parent_t, Args...>(hash, name, value, inputs, eInfos);
   }
 
   /// 3. generic case
   template <typename R, typename C, soa::is_table... Args>
-  static void inputsFromArgs(R (C::*)(Args...), const char* name, bool value, std::vector<InputSpec>& inputs, std::vector<ExpressionInfo>& eInfos, std::vector<StringPair>&, std::vector<StringPair>&)
+  static void inputsFromArgs(R (C::*)(Args...), const char* name, bool value, std::vector<InputSpec>& inputs, std::vector<ExpressionInfo>& eInfos, Cache&, Cache&)
     requires(std::is_lvalue_reference_v<Args> && ...)
   {
     constexpr auto hash = o2::framework::TypeIdHelpers::uniqueId<R (C::*)(Args...)>();
@@ -480,8 +480,8 @@ DataProcessorSpec adaptAnalysisTask(ConfigContext const& ctx, Args&&... args)
   std::vector<InputSpec> inputs;
   std::vector<ConfigParamSpec> options;
   std::vector<ExpressionInfo> expressionInfos;
-  std::vector<StringPair> bindingsKeys;
-  std::vector<StringPair> bindingsKeysUnsorted;
+  Cache bindingsKeys;
+  Cache bindingsKeysUnsorted;
 
   /// make sure options and configurables are set before expression infos are created
   homogeneous_apply_refs([&options, &hash](auto& element) { return analysis_task_parsers::appendOption(options, element); }, *task.get());

Framework/Core/include/Framework/ArrowTableSlicingCache.h

@@ -34,51 +34,64 @@ struct SliceInfoUnsortedPtr {
   gsl::span<int64_t const> getSliceFor(int value) const;
 };
 
-using StringPair = std::pair<std::string, std::string>;
+struct Entry {
+  std::string binding;
+  std::string key;
+  bool enabled;
+
+  Entry(std::string b, std::string k, bool e = true)
+    : binding{b},
+      key{k},
+      enabled{e}
+  {
+  }
+};
+
+using Cache = std::vector<Entry>;
 
-void updatePairList(std::vector<StringPair>& list, std::string const& binding, std::string const& key);
+void updatePairList(Cache& list, std::string const& binding, std::string const& key, bool enabled);
 
 struct ArrowTableSlicingCacheDef {
   constexpr static ServiceKind service_kind = ServiceKind::Global;
-  std::vector<StringPair> bindingsKeys;
-  std::vector<StringPair> bindingsKeysUnsorted;
+  Cache bindingsKeys;
+  Cache bindingsKeysUnsorted;
 
-  void setCaches(std::vector<StringPair>&& bsks);
-  void setCachesUnsorted(std::vector<StringPair>&& bsks);
+  void setCaches(Cache&& bsks);
+  void setCachesUnsorted(Cache&& bsks);
 };
 
 struct ArrowTableSlicingCache {
   constexpr static ServiceKind service_kind = ServiceKind::Stream;
 
-  std::vector<StringPair> bindingsKeys;
+  Cache bindingsKeys;
   std::vector<std::shared_ptr<arrow::NumericArray<arrow::Int32Type>>> values;
   std::vector<std::shared_ptr<arrow::NumericArray<arrow::Int64Type>>> counts;
 
-  std::vector<StringPair> bindingsKeysUnsorted;
+  Cache bindingsKeysUnsorted;
   std::vector<std::vector<int>> valuesUnsorted;
   std::vector<ListVector> groups;
 
-  ArrowTableSlicingCache(std::vector<StringPair>&& bsks, std::vector<StringPair>&& bsksUnsorted = {});
+  ArrowTableSlicingCache(Cache&& bsks, Cache&& bsksUnsorted = {});
 
   // set caching information externally
-  void setCaches(std::vector<StringPair>&& bsks, std::vector<StringPair>&& bsksUnsorted = {});
+  void setCaches(Cache&& bsks, Cache&& bsksUnsorted = {});
 
   // update slicing info cache entry (assumes it is already present)
   arrow::Status updateCacheEntry(int pos, std::shared_ptr<arrow::Table> const& table);
   arrow::Status updateCacheEntryUnsorted(int pos, std::shared_ptr<arrow::Table> const& table);
 
   // helper to locate cache position
-  std::pair<int, bool> getCachePos(StringPair const& bindingKey) const;
-  int getCachePosSortedFor(StringPair const& bindingKey) const;
-  int getCachePosUnsortedFor(StringPair const& bindingKey) const;
+  std::pair<int, bool> getCachePos(Entry const& bindingKey) const;
+  int getCachePosSortedFor(Entry const& bindingKey) const;
+  int getCachePosUnsortedFor(Entry const& bindingKey) const;
 
   // get slice from cache for a given value
-  SliceInfoPtr getCacheFor(StringPair const& bindingKey) const;
-  SliceInfoUnsortedPtr getCacheUnsortedFor(StringPair const& bindingKey) const;
+  SliceInfoPtr getCacheFor(Entry const& bindingKey) const;
+  SliceInfoUnsortedPtr getCacheUnsortedFor(Entry const& bindingKey) const;
   SliceInfoPtr getCacheForPos(int pos) const;
   SliceInfoUnsortedPtr getCacheUnsortedForPos(int pos) const;
 
-  static void validateOrder(StringPair const& bindingKey, std::shared_ptr<arrow::Table> const& input);
+  static void validateOrder(Entry const& bindingKey, std::shared_ptr<arrow::Table> const& input);
 };
 } // namespace o2::framework
 

Framework/Core/include/Framework/GroupSlicer.h

@@ -55,7 +55,7 @@ struct GroupSlicer {
     {
       constexpr auto index = framework::has_type_at_v<std::decay_t<T>>(associated_pack_t{});
       auto binding = o2::soa::getLabelFromTypeForKey<std::decay_t<T>>(mIndexColumnName);
-      auto bk = std::make_pair(binding, mIndexColumnName);
+      auto bk = Entry(binding, mIndexColumnName);
       if constexpr (!o2::soa::is_smallgroups<std::decay_t<T>>) {
         if (table.size() == 0) {
           return;

Framework/Core/src/ASoA.cxx

@@ -197,7 +197,7 @@ bool PreslicePolicyBase::isMissing() const
   return binding == "[MISSING]";
 }
 
-StringPair const& PreslicePolicyBase::getBindingKey() const
+Entry const& PreslicePolicyBase::getBindingKey() const
 {
   return bindingKey;
 }

Framework/Core/src/ArrowSupport.cxx

@@ -567,26 +567,27 @@ o2::framework::ServiceSpec ArrowSupport::arrowTableSlicingCacheSpec()
     .name = "arrow-slicing-cache",
     .uniqueId = CommonServices::simpleServiceId<ArrowTableSlicingCache>(),
     .init = [](ServiceRegistryRef services, DeviceState&, fair::mq::ProgOptions&) { return ServiceHandle{TypeIdHelpers::uniqueId<ArrowTableSlicingCache>(),
-                                                                                                         new ArrowTableSlicingCache(std::vector<std::pair<std::string, std::string>>{services.get<ArrowTableSlicingCacheDef>().bindingsKeys}, std::vector{services.get<ArrowTableSlicingCacheDef>().bindingsKeysUnsorted}),
+                                                                                                         new ArrowTableSlicingCache(Cache{services.get<ArrowTableSlicingCacheDef>().bindingsKeys},
+                                                                                                                                    Cache{services.get<ArrowTableSlicingCacheDef>().bindingsKeysUnsorted}),
                                                                                                          ServiceKind::Stream, typeid(ArrowTableSlicingCache).name()}; },
     .configure = CommonServices::noConfiguration(),
     .preProcessing = [](ProcessingContext& pc, void* service_ptr) {
       auto* service = static_cast<ArrowTableSlicingCache*>(service_ptr);
       auto& caches = service->bindingsKeys;
-      for (auto i = 0; i < caches.size(); ++i) {
-        if (pc.inputs().getPos(caches[i].first.c_str()) >= 0) {
-          auto status = service->updateCacheEntry(i, pc.inputs().get<TableConsumer>(caches[i].first.c_str())->asArrowTable());
+      for (auto i = 0u; i < caches.size(); ++i) {
+        if (caches[i].enabled && pc.inputs().getPos(caches[i].binding.c_str()) >= 0) {
+          auto status = service->updateCacheEntry(i, pc.inputs().get<TableConsumer>(caches[i].binding.c_str())->asArrowTable());
           if (!status.ok()) {
-            throw runtime_error_f("Failed to update slice cache for %s/%s", caches[i].first.c_str(), caches[i].second.c_str());
+            throw runtime_error_f("Failed to update slice cache for %s/%s", caches[i].binding.c_str(), caches[i].key.c_str());
           }
         }
       }
       auto& unsortedCaches = service->bindingsKeysUnsorted;
-      for (auto i = 0; i < unsortedCaches.size(); ++i) {
-        if (pc.inputs().getPos(unsortedCaches[i].first.c_str()) >= 0) {
-          auto status = service->updateCacheEntryUnsorted(i, pc.inputs().get<TableConsumer>(unsortedCaches[i].first.c_str())->asArrowTable());
+      for (auto i = 0u; i < unsortedCaches.size(); ++i) {
+        if (unsortedCaches[i].enabled && pc.inputs().getPos(unsortedCaches[i].binding.c_str()) >= 0) {
+          auto status = service->updateCacheEntryUnsorted(i, pc.inputs().get<TableConsumer>(unsortedCaches[i].binding.c_str())->asArrowTable());
           if (!status.ok()) {
-            throw runtime_error_f("failed to update slice cache (unsorted) for %s/%s", unsortedCaches[i].first.c_str(), unsortedCaches[i].second.c_str());
+            throw runtime_error_f("failed to update slice cache (unsorted) for %s/%s", unsortedCaches[i].binding.c_str(), unsortedCaches[i].key.c_str());
           }
         }
       } },