ceph/src/arrow/cpp/src/gandiva/expr_decomposer.h

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements.  See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership.  The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License.  You may obtain a copy of the License at
   8 //
   9 //   http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing,
  12 // software distributed under the License is distributed on an
  13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 // KIND, either express or implied.  See the License for the
  15 // specific language governing permissions and limitations
  16 // under the License.
  17
  18 #pragma once
  19
  20 #include <cmath>
  21 #include <memory>
  22 #include <stack>
  23 #include <string>
  24 #include <utility>
  25
  26 #include "gandiva/arrow.h"
  27 #include "gandiva/expression.h"
  28 #include "gandiva/node.h"
  29 #include "gandiva/node_visitor.h"
  30 #include "gandiva/visibility.h"
  31
  32 namespace gandiva {
  33
  34 class FunctionRegistry;
  35 class Annotator;
  36
  37 /// \brief Decomposes an expression tree to separate out the validity and
  38 /// value expressions.
  39 class GANDIVA_EXPORT ExprDecomposer : public NodeVisitor {
  40  public:
  41   explicit ExprDecomposer(const FunctionRegistry& registry, Annotator& annotator)
  42       : registry_(registry), annotator_(annotator), nested_if_else_(false) {}
  43
  44   Status Decompose(const Node& root, ValueValidityPairPtr* out) {
  45     auto status = root.Accept(*this);
  46     if (status.ok()) {
  47       *out = std::move(result_);
  48     }
  49     return status;
  50   }
  51
  52  private:
  53   ARROW_DISALLOW_COPY_AND_ASSIGN(ExprDecomposer);
  54
  55   FRIEND_TEST(TestExprDecomposer, TestStackSimple);
  56   FRIEND_TEST(TestExprDecomposer, TestNested);
  57   FRIEND_TEST(TestExprDecomposer, TestInternalIf);
  58   FRIEND_TEST(TestExprDecomposer, TestParallelIf);
  59   FRIEND_TEST(TestExprDecomposer, TestIfInCondition);
  60   FRIEND_TEST(TestExprDecomposer, TestFunctionBetweenNestedIf);
  61   FRIEND_TEST(TestExprDecomposer, TestComplexIfCondition);
  62
  63   Status Visit(const FieldNode& node) override;
  64   Status Visit(const FunctionNode& node) override;
  65   Status Visit(const IfNode& node) override;
  66   Status Visit(const LiteralNode& node) override;
  67   Status Visit(const BooleanNode& node) override;
  68   Status Visit(const InExpressionNode<int32_t>& node) override;
  69   Status Visit(const InExpressionNode<int64_t>& node) override;
  70   Status Visit(const InExpressionNode<float>& node) override;
  71   Status Visit(const InExpressionNode<double>& node) override;
  72   Status Visit(const InExpressionNode<gandiva::DecimalScalar128>& node) override;
  73   Status Visit(const InExpressionNode<std::string>& node) override;
  74
  75   // Optimize a function node, if possible.
  76   const FunctionNode TryOptimize(const FunctionNode& node);
  77
  78   enum StackEntryType { kStackEntryCondition, kStackEntryThen, kStackEntryElse };
  79
  80   // stack of if nodes.
  81   class IfStackEntry {
  82    public:
  83     IfStackEntry(const IfNode& if_node, StackEntryType entry_type,
  84                  bool is_terminal_else = false, int local_bitmap_idx = 0)
  85         : if_node_(if_node),
  86           entry_type_(entry_type),
  87           is_terminal_else_(is_terminal_else),
  88           local_bitmap_idx_(local_bitmap_idx) {}
  89
  90     const IfNode& if_node_;
  91     StackEntryType entry_type_;
  92     bool is_terminal_else_;
  93     int local_bitmap_idx_;
  94
  95    private:
  96     ARROW_DISALLOW_COPY_AND_ASSIGN(IfStackEntry);
  97   };
  98
  99   // pop 'condition entry' into stack.
 100   void PushConditionEntry(const IfNode& node);
 101
 102   // pop 'condition entry' from stack.
 103   void PopConditionEntry(const IfNode& node);
 104
 105   // push 'then entry' to stack. returns either a new local bitmap or the parent's
 106   // bitmap (in case of nested if-else).
 107   int PushThenEntry(const IfNode& node, bool reuse_bitmap);
 108
 109   // pop 'then entry' from stack.
 110   void PopThenEntry(const IfNode& node);
 111
 112   // push 'else entry' into stack.
 113   void PushElseEntry(const IfNode& node, int local_bitmap_idx);
 114
 115   // pop 'else entry' from stack. returns 'true' if this is a terminal else condition
 116   // i.e no nested if condition below this node.
 117   bool PopElseEntry(const IfNode& node);
 118
 119   ValueValidityPairPtr result() { return std::move(result_); }
 120
 121   const FunctionRegistry& registry_;
 122   Annotator& annotator_;
 123   std::stack<std::unique_ptr<IfStackEntry>> if_entries_stack_;
 124   ValueValidityPairPtr result_;
 125   bool nested_if_else_;
 126 };
 127
 128 }  // namespace gandiva