ceph/src/arrow/cpp/src/gandiva/tests/literal_test.cc

   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 #include <gtest/gtest.h>
  19 #include "arrow/memory_pool.h"
  20 #include "arrow/status.h"
  21
  22 #include "gandiva/projector.h"
  23 #include "gandiva/tests/test_util.h"
  24 #include "gandiva/tree_expr_builder.h"
  25
  26 namespace gandiva {
  27
  28 using arrow::boolean;
  29 using arrow::float32;
  30 using arrow::float64;
  31 using arrow::int32;
  32 using arrow::int64;
  33
  34 class TestLiteral : public ::testing::Test {
  35  public:
  36   void SetUp() { pool_ = arrow::default_memory_pool(); }
  37
  38  protected:
  39   arrow::MemoryPool* pool_;
  40 };
  41
  42 TEST_F(TestLiteral, TestSimpleArithmetic) {
  43   // schema for input fields
  44   auto field_a = field("a", boolean());
  45   auto field_b = field("b", int32());
  46   auto field_c = field("c", int64());
  47   auto field_d = field("d", float32());
  48   auto field_e = field("e", float64());
  49   auto schema = arrow::schema({field_a, field_b, field_c, field_d, field_e});
  50
  51   // output fields
  52   auto res_a = field("a+1", boolean());
  53   auto res_b = field("b+1", int32());
  54   auto res_c = field("c+1", int64());
  55   auto res_d = field("d+1", float32());
  56   auto res_e = field("e+1", float64());
  57
  58   // build expressions.
  59   // a == true
  60   // b + 1
  61   // c + 1
  62   // d + 1
  63   // e + 1
  64   auto node_a = TreeExprBuilder::MakeField(field_a);
  65   auto literal_a = TreeExprBuilder::MakeLiteral(true);
  66   auto func_a = TreeExprBuilder::MakeFunction("equal", {node_a, literal_a}, boolean());
  67   auto expr_a = TreeExprBuilder::MakeExpression(func_a, res_a);
  68
  69   auto node_b = TreeExprBuilder::MakeField(field_b);
  70   auto literal_b = TreeExprBuilder::MakeLiteral((int32_t)1);
  71   auto func_b = TreeExprBuilder::MakeFunction("add", {node_b, literal_b}, int32());
  72   auto expr_b = TreeExprBuilder::MakeExpression(func_b, res_b);
  73
  74   auto node_c = TreeExprBuilder::MakeField(field_c);
  75   auto literal_c = TreeExprBuilder::MakeLiteral((int64_t)1);
  76   auto func_c = TreeExprBuilder::MakeFunction("add", {node_c, literal_c}, int64());
  77   auto expr_c = TreeExprBuilder::MakeExpression(func_c, res_c);
  78
  79   auto node_d = TreeExprBuilder::MakeField(field_d);
  80   auto literal_d = TreeExprBuilder::MakeLiteral(static_cast<float>(1));
  81   auto func_d = TreeExprBuilder::MakeFunction("add", {node_d, literal_d}, float32());
  82   auto expr_d = TreeExprBuilder::MakeExpression(func_d, res_d);
  83
  84   auto node_e = TreeExprBuilder::MakeField(field_e);
  85   auto literal_e = TreeExprBuilder::MakeLiteral(static_cast<double>(1));
  86   auto func_e = TreeExprBuilder::MakeFunction("add", {node_e, literal_e}, float64());
  87   auto expr_e = TreeExprBuilder::MakeExpression(func_e, res_e);
  88
  89   // Build a projector for the expressions.
  90   std::shared_ptr<Projector> projector;
  91   auto status = Projector::Make(schema, {expr_a, expr_b, expr_c, expr_d, expr_e},
  92                                 TestConfiguration(), &projector);
  93   EXPECT_TRUE(status.ok());
  94
  95   // Create a row-batch with some sample data
  96   int num_records = 4;
  97   auto array_a = MakeArrowArrayBool({true, true, false, true}, {true, true, true, false});
  98   auto array_b = MakeArrowArrayInt32({5, 15, -15, 17}, {true, true, true, false});
  99   auto array_c = MakeArrowArrayInt64({5, 15, -15, 17}, {true, true, true, false});
 100   auto array_d = MakeArrowArrayFloat32({5.2f, 15, -15.6f, 17}, {true, true, true, false});
 101   auto array_e = MakeArrowArrayFloat64({5.6f, 15, -15.9f, 17}, {true, true, true, false});
 102
 103   // expected output
 104   auto exp_a = MakeArrowArrayBool({true, true, false, false}, {true, true, true, false});
 105   auto exp_b = MakeArrowArrayInt32({6, 16, -14, 0}, {true, true, true, false});
 106   auto exp_c = MakeArrowArrayInt64({6, 16, -14, 0}, {true, true, true, false});
 107   auto exp_d = MakeArrowArrayFloat32({6.2f, 16, -14.6f, 0}, {true, true, true, false});
 108   auto exp_e = MakeArrowArrayFloat64({6.6f, 16, -14.9f, 0}, {true, true, true, false});
 109
 110   // prepare input record batch
 111   auto in_batch = arrow::RecordBatch::Make(schema, num_records,
 112                                            {array_a, array_b, array_c, array_d, array_e});
 113
 114   // Evaluate expression
 115   arrow::ArrayVector outputs;
 116   status = projector->Evaluate(*in_batch, pool_, &outputs);
 117   EXPECT_TRUE(status.ok());
 118
 119   // Validate results
 120   EXPECT_ARROW_ARRAY_EQUALS(exp_a, outputs.at(0));
 121   EXPECT_ARROW_ARRAY_EQUALS(exp_b, outputs.at(1));
 122   EXPECT_ARROW_ARRAY_EQUALS(exp_c, outputs.at(2));
 123   EXPECT_ARROW_ARRAY_EQUALS(exp_d, outputs.at(3));
 124   EXPECT_ARROW_ARRAY_EQUALS(exp_e, outputs.at(4));
 125 }
 126
 127 TEST_F(TestLiteral, TestLiteralHash) {
 128   auto schema = arrow::schema({});
 129   // output fields
 130   auto res = field("a", int32());
 131   auto int_literal = TreeExprBuilder::MakeLiteral((int32_t)2);
 132   auto expr = TreeExprBuilder::MakeExpression(int_literal, res);
 133
 134   // Build a projector for the expressions.
 135   std::shared_ptr<Projector> projector;
 136   auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
 137   EXPECT_TRUE(status.ok()) << status.message();
 138
 139   auto res1 = field("a", int64());
 140   auto int_literal1 = TreeExprBuilder::MakeLiteral((int64_t)2);
 141   auto expr1 = TreeExprBuilder::MakeExpression(int_literal1, res1);
 142
 143   // Build a projector for the expressions.
 144   std::shared_ptr<Projector> projector1;
 145   status = Projector::Make(schema, {expr1}, TestConfiguration(), &projector1);
 146   EXPECT_TRUE(status.ok()) << status.message();
 147   EXPECT_TRUE(projector.get() != projector1.get());
 148 }
 149
 150 TEST_F(TestLiteral, TestNullLiteral) {
 151   // schema for input fields
 152   auto field_a = field("a", int32());
 153   auto field_b = field("b", int32());
 154   auto schema = arrow::schema({field_a, field_b});
 155
 156   // output fields
 157   auto res = field("a+b+null", int32());
 158
 159   auto node_a = TreeExprBuilder::MakeField(field_a);
 160   auto node_b = TreeExprBuilder::MakeField(field_b);
 161   auto literal_c = TreeExprBuilder::MakeNull(arrow::int32());
 162   auto add_a_b = TreeExprBuilder::MakeFunction("add", {node_a, node_b}, int32());
 163   auto add_a_b_c = TreeExprBuilder::MakeFunction("add", {add_a_b, literal_c}, int32());
 164   auto expr = TreeExprBuilder::MakeExpression(add_a_b_c, res);
 165
 166   // Build a projector for the expressions.
 167   std::shared_ptr<Projector> projector;
 168   auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
 169   EXPECT_TRUE(status.ok()) << status.message();
 170
 171   // Create a row-batch with some sample data
 172   int num_records = 4;
 173   auto array_a = MakeArrowArrayInt32({5, 15, -15, 17}, {true, true, true, false});
 174   auto array_b = MakeArrowArrayInt32({5, 15, -15, 17}, {true, true, true, false});
 175
 176   // expected output
 177   auto exp = MakeArrowArrayInt32({0, 0, 0, 0}, {false, false, false, false});
 178
 179   // prepare input record batch
 180   auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array_a, array_b});
 181
 182   // Evaluate expression
 183   arrow::ArrayVector outputs;
 184   status = projector->Evaluate(*in_batch, pool_, &outputs);
 185   EXPECT_TRUE(status.ok()) << status.message();
 186
 187   // Validate results
 188   EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
 189 }
 190
 191 TEST_F(TestLiteral, TestNullLiteralInIf) {
 192   // schema for input fields
 193   auto field_a = field("a", float64());
 194   auto schema = arrow::schema({field_a});
 195
 196   // output fields
 197   auto res = field("res", float64());
 198
 199   auto node_a = TreeExprBuilder::MakeField(field_a);
 200   auto literal_5 = TreeExprBuilder::MakeLiteral(5.0);
 201   auto a_gt_5 = TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_5},
 202                                               arrow::boolean());
 203   auto literal_null = TreeExprBuilder::MakeNull(arrow::float64());
 204   auto if_node =
 205       TreeExprBuilder::MakeIf(a_gt_5, literal_5, literal_null, arrow::float64());
 206   auto expr = TreeExprBuilder::MakeExpression(if_node, res);
 207
 208   // Build a projector for the expressions.
 209   std::shared_ptr<Projector> projector;
 210   auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
 211   EXPECT_TRUE(status.ok()) << status.message();
 212
 213   // Create a row-batch with some sample data
 214   int num_records = 4;
 215   auto array_a = MakeArrowArrayFloat64({6, 15, -15, 17}, {true, true, true, false});
 216
 217   // expected output
 218   auto exp = MakeArrowArrayFloat64({5, 5, 0, 0}, {true, true, false, false});
 219
 220   // prepare input record batch
 221   auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array_a});
 222
 223   // Evaluate expression
 224   arrow::ArrayVector outputs;
 225   status = projector->Evaluate(*in_batch, pool_, &outputs);
 226   EXPECT_TRUE(status.ok()) << status.message();
 227
 228   // Validate results
 229   EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
 230 }
 231
 232 }  // namespace gandiva