DPL Analysis: Introducing conditional expressions (#6753)

`ifnode(condition, then, else)` operation is added to expressions. These can be nested and all three arguments can be arbitrary valid expressions. `condition` needs to have boolean result, `then` and `else` should return similar types (ideally the same - both floats, or both boolean, etc.). A `conditionalExpressions.cxx` tutorial example is added (note that it uses bitwise operations in filter expression and thus will only work as is with arrow > 3).

Author: aalkin

Analysis/Tutorials/CMakeLists.txt

@@ -228,3 +228,9 @@ o2_add_dpl_workflow(multiprocess-example
                   JOB_POOL analysis
                   PUBLIC_LINK_LIBRARIES O2::Framework O2::AnalysisCore O2::AnalysisDataModel
                   COMPONENT_NAME AnalysisTutorial)
+
+o2_add_dpl_workflow(conditional-expressions
+                  SOURCES src/conditionalExpressions.cxx
+                  JOB_POOL analysis
+                  PUBLIC_LINK_LIBRARIES O2::Framework O2::AnalysisCore O2::AnalysisDataModel
+                  COMPONENT_NAME AnalysisTutorial)

Analysis/Tutorials/src/conditionalExpressions.cxx

@@ -0,0 +1,56 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// 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.
+///
+/// \brief Demonstration of conditions in filter expressions
+
+#include "Framework/runDataProcessing.h"
+#include "Framework/AnalysisTask.h"
+
+using namespace o2;
+using namespace o2::framework;
+using namespace o2::framework::expressions;
+
+struct ConditionalExpressions {
+  Configurable<bool> useFlags{"useFlags", false, "Switch to enable using track flags for selection"};
+  Filter trackFilter = nabs(aod::track::eta) < 0.9f && aod::track::pt > 0.5f && ifnode(useFlags.node() == true, (aod::track::flags & static_cast<uint32_t>(o2::aod::track::ITSrefit)) != 0u, true);
+  OutputObj<TH2F> etapt{TH2F("etapt", ";#eta;#p_{T}", 201, -2.1, 2.1, 601, 0, 60.1)};
+  void process(aod::Collision const&, soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra>> const& tracks)
+  {
+    for (auto& track : tracks) {
+      etapt->Fill(track.eta(), track.pt());
+    }
+  }
+};
+
+struct BasicOperations {
+  Configurable<bool> useFlags{"useFlags", false, "Switch to enable using track flags for selection"};
+  Filter trackFilter = nabs(aod::track::eta) < 0.9f && aod::track::pt > 0.5f;
+  OutputObj<TH2F> etapt{TH2F("etapt", ";#eta;#p_{T}", 201, -2.1, 2.1, 601, 0, 60.1)};
+  void process(aod::Collision const&, soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra>> const& tracks)
+  {
+    for (auto& track : tracks) {
+      if (useFlags) {
+        if ((track.flags() & o2::aod::track::ITSrefit) != 0u) {
+          etapt->Fill(track.eta(), track.pt());
+        }
+      } else {
+        etapt->Fill(track.eta(), track.pt());
+      }
+    }
+  }
+};
+
+WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
+{
+  return WorkflowSpec{
+    adaptAnalysisTask<ConditionalExpressions>(cfgc),
+    adaptAnalysisTask<BasicOperations>(cfgc)};
+}

Framework/Core/include/Framework/BasicOps.h

@@ -41,7 +41,8 @@ enum BasicOp : unsigned int {
   Acos,
   Atan,
   Abs,
-  BitwiseNot
+  BitwiseNot,
+  Conditional
 };
 } // namespace o2::framework
 

Framework/Core/include/Framework/Configurable.h

@@ -15,6 +15,11 @@
 #include <vector>
 namespace o2::framework
 {
+namespace expressions
+{
+struct PlaceholderNode;
+}
+
 template <typename T, ConfigParamKind K>
 struct ConfigurableBase {
   ConfigurableBase(std::string const& name, T&& defaultValue, std::string const& help)
@@ -68,6 +73,10 @@ struct Configurable : IP {
     : IP{name, std::forward<T>(defaultValue), help}
   {
   }
+  auto node()
+  {
+    return expressions::PlaceholderNode{*this};
+  }
 };
 
 template <typename T, ConfigParamKind K = ConfigParamKind::kGeneric>

Framework/Core/include/Framework/ExpressionHelpers.h

@@ -20,7 +20,7 @@ namespace o2::framework::expressions
 {
 /// a map between BasicOp and gandiva node definitions
 /// note that logical 'and' and 'or' are created separately
-static std::array<std::string, BasicOp::BitwiseNot + 1> basicOperationsMap = {
+static std::array<std::string, BasicOp::Conditional + 1> basicOperationsMap = {
   "and",
   "or",
   "add",
@@ -48,7 +48,8 @@ static std::array<std::string, BasicOp::BitwiseNot + 1> basicOperationsMap = {
   "acosf",
   "atanf",
   "absf",
-  "bitwise_not"};
+  "bitwise_not",
+  "if"};
 
 struct DatumSpec {
   /// datum spec either contains an index, a value of a literal or a binding label
@@ -72,17 +73,21 @@ bool operator==(DatumSpec const& lhs, DatumSpec const& rhs);
 std::ostream& operator<<(std::ostream& os, DatumSpec const& spec);
 
 struct ColumnOperationSpec {
+  size_t index = 0;
   BasicOp op;
   DatumSpec left;
   DatumSpec right;
+  DatumSpec condition;
   DatumSpec result;
   atype::type type = atype::NA;
   ColumnOperationSpec() = default;
-  // TODO: extend this to support unary ops seamlessly
-  explicit ColumnOperationSpec(BasicOp op_) : op{op_},
-                                              left{},
-                                              right{},
-                                              result{}
+  explicit ColumnOperationSpec(BasicOp op_, size_t index_ = 0)
+    : index{index_},
+      op{op_},
+      left{},
+      right{},
+      condition{},
+      result{}
   {
     switch (op) {
       case BasicOp::LogicalOr:
@@ -110,6 +115,10 @@ struct NodeRecord {
   Node* node_ptr = nullptr;
   size_t index = 0;
   explicit NodeRecord(Node* node_, size_t index_) : node_ptr(node_), index{index_} {}
+  bool operator!=(NodeRecord const& rhs)
+  {
+    return this->node_ptr != rhs.node_ptr;
+  }
 };
 } // namespace o2::framework::expressions
 

Framework/Core/include/Framework/Expressions.h

@@ -128,7 +128,7 @@ struct OpNode {
 /// A placeholder node for simple type configurable
 struct PlaceholderNode : LiteralNode {
   template <typename T>
-  PlaceholderNode(Configurable<T> v) : LiteralNode{v.value}, name{v.name}
+  PlaceholderNode(Configurable<T> const& v) : LiteralNode{v.value}, name{v.name}
   {
     if constexpr (variant_trait_v<typename std::decay<T>::type> != VariantType::Unknown) {
       retrieve = [](InitContext& context, std::string const& name) { return LiteralNode::var_t{context.options().get<T>(name.c_str())}; };
@@ -146,40 +146,54 @@ struct PlaceholderNode : LiteralNode {
   LiteralNode::var_t (*retrieve)(InitContext&, std::string const& name);
 };
 
+/// A conditional node
+struct ConditionalNode {
+};
+
 /// A generic tree node
 struct Node {
-  Node(LiteralNode v) : self{v}, left{nullptr}, right{nullptr}
+  Node(LiteralNode v) : self{v}, left{nullptr}, right{nullptr}, condition{nullptr}
   {
   }
 
-  Node(PlaceholderNode v) : self{v}, left{nullptr}, right{nullptr}
+  Node(PlaceholderNode v) : self{v}, left{nullptr}, right{nullptr}, condition{nullptr}
   {
   }
 
-  Node(Node&& n) : self{n.self}, left{std::move(n.left)}, right{std::move(n.right)}
+  Node(Node&& n) : self{n.self}, left{std::move(n.left)}, right{std::move(n.right)}, condition{std::move(n.condition)}
   {
   }
 
-  Node(BindingNode n) : self{n}, left{nullptr}, right{nullptr}
+  Node(BindingNode n) : self{n}, left{nullptr}, right{nullptr}, condition{nullptr}
   {
   }
 
+  Node(ConditionalNode op, Node&& then_, Node&& else_, Node&& condition_)
+    : self{op},
+      left{std::make_unique<Node>(std::move(then_))},
+      right{std::make_unique<Node>(std::move(else_))},
+      condition{std::make_unique<Node>(std::move(condition_))} {}
+
   Node(OpNode op, Node&& l, Node&& r)
     : self{op},
       left{std::make_unique<Node>(std::move(l))},
-      right{std::make_unique<Node>(std::move(r))} {}
+      right{std::make_unique<Node>(std::move(r))},
+      condition{nullptr} {}
 
   Node(OpNode op, Node&& l)
     : self{op},
       left{std::make_unique<Node>(std::move(l))},
-      right{nullptr} {}
+      right{nullptr},
+      condition{nullptr} {}
 
   /// variant with possible nodes
-  using self_t = std::variant<LiteralNode, BindingNode, OpNode, PlaceholderNode>;
+  using self_t = std::variant<LiteralNode, BindingNode, OpNode, PlaceholderNode, ConditionalNode>;
   self_t self;
+  size_t index = 0;
   /// pointers to children
   std::unique_ptr<Node> left;
   std::unique_ptr<Node> right;
+  std::unique_ptr<Node> condition;
 };
 
 /// overloaded operators to build the tree from an expression
@@ -319,20 +333,84 @@ inline Node nbitwise_not(Node left)
   return Node{OpNode{BasicOp::BitwiseNot}, std::move(left)};
 }
 
+/// conditionals
+template <typename C, typename T, typename E>
+inline Node ifnode(C condition_, T then_, E else_)
+{
+  return Node{ConditionalNode{}, std::move(then_), std::move(else_), std::move(condition_)};
+}
+
+template <>
+inline Node ifnode(Node condition_, Node then_, Node else_)
+{
+  return Node{ConditionalNode{}, std::move(then_), std::move(else_), std::move(condition_)};
+}
+
+template <typename L, std::enable_if_t<std::is_integral<L>::value || std::is_floating_point<L>::value, bool> = true>
+inline Node ifnode(Node condition_, Node then_, L else_)
+{
+  return Node{ConditionalNode{}, std::move(then_), LiteralNode{else_}, std::move(condition_)};
+}
+
+template <typename L, std::enable_if_t<std::is_integral<L>::value || std::is_floating_point<L>::value, bool> = true>
+inline Node ifnode(Node condition_, L then_, Node else_)
+{
+  return Node{ConditionalNode{}, LiteralNode{then_}, std::move(else_), std::move(condition_)};
+}
+
+template <typename L1, typename L2, std::enable_if_t<(std::is_integral<L1>::value || std::is_floating_point<L1>::value) && (std::is_integral<L2>::value || std::is_floating_point<L2>::value), bool> = true>
+inline Node ifnode(Node condition_, L1 then_, L2 else_)
+{
+  return Node{ConditionalNode{}, LiteralNode{then_}, LiteralNode{else_}, std::move(condition_)};
+}
+
+template <typename T>
+inline Node ifnode(Configurable<T> condition_, Node then_, Node else_)
+{
+  return Node{ConditionalNode{}, std::move(then_), std::move(else_), PlaceholderNode{condition_}};
+}
+
+template <typename L>
+inline Node ifnode(Node condition_, Node then_, Configurable<L> else_)
+{
+  return Node{ConditionalNode{}, std::move(then_), PlaceholderNode{else_}, std::move(condition_)};
+}
+
+template <typename L>
+inline Node ifnode(Node condition_, Configurable<L> then_, Node else_)
+{
+  return Node{ConditionalNode{}, PlaceholderNode{then_}, std::move(else_), std::move(condition_)};
+}
+
+template <typename L1, typename L2>
+inline Node ifnode(Node condition_, Configurable<L1> then_, Configurable<L2> else_)
+{
+  return Node{ConditionalNode{}, PlaceholderNode{then_}, PlaceholderNode{else_}, std::move(condition_)};
+}
+
 /// A struct, containing the root of the expression tree
 struct Filter {
-  Filter(Node&& node_) : node{std::make_unique<Node>(std::move(node_))} {}
-  Filter(Filter&& other) : node{std::move(other.node)} {}
+  Filter(Node&& node_) : node{std::make_unique<Node>(std::move(node_))}
+  {
+    (void)designateSubtrees(node.get());
+  }
+
+  Filter(Filter&& other) : node{std::move(other.node)}
+  {
+    (void)designateSubtrees(node.get());
+  }
   std::unique_ptr<Node> node;
+
+  size_t designateSubtrees(Node* node, size_t index = 0);
 };
 
 using Projector = Filter;
 
 using Selection = std::shared_ptr<gandiva::SelectionVector>;
 /// Function for creating gandiva selection from our internal filter tree
-Selection createSelection(std::shared_ptr<arrow::Table> table, Filter const& expression);
+Selection createSelection(std::shared_ptr<arrow::Table> const& table, Filter const& expression);
 /// Function for creating gandiva selection from prepared gandiva expressions tree
-Selection createSelection(std::shared_ptr<arrow::Table> table, std::shared_ptr<gandiva::Filter> gfilter);
+Selection createSelection(std::shared_ptr<arrow::Table> const& table, std::shared_ptr<gandiva::Filter> gfilter);
 
 struct ColumnOperationSpec;
 using Operations = std::vector<ColumnOperationSpec>;