]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 |