]>
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::int32; | |
30 | ||
31 | class TestBooleanExpr : public ::testing::Test { | |
32 | public: | |
33 | void SetUp() { pool_ = arrow::default_memory_pool(); } | |
34 | ||
35 | protected: | |
36 | arrow::MemoryPool* pool_; | |
37 | }; | |
38 | ||
39 | TEST_F(TestBooleanExpr, SimpleAnd) { | |
40 | // schema for input fields | |
41 | auto fielda = field("a", int32()); | |
42 | auto fieldb = field("b", int32()); | |
43 | auto schema = arrow::schema({fielda, fieldb}); | |
44 | ||
45 | // output fields | |
46 | auto field_result = field("res", boolean()); | |
47 | ||
48 | // build expression. | |
49 | // (a > 0) && (b > 0) | |
50 | auto node_a = TreeExprBuilder::MakeField(fielda); | |
51 | auto node_b = TreeExprBuilder::MakeField(fieldb); | |
52 | auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0); | |
53 | auto a_gt_0 = | |
54 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean()); | |
55 | auto b_gt_0 = | |
56 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean()); | |
57 | ||
58 | auto node_and = TreeExprBuilder::MakeAnd({a_gt_0, b_gt_0}); | |
59 | auto expr = TreeExprBuilder::MakeExpression(node_and, field_result); | |
60 | ||
61 | // Build a projector for the expressions. | |
62 | std::shared_ptr<Projector> projector; | |
63 | auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector); | |
64 | EXPECT_TRUE(status.ok()); | |
65 | ||
66 | // FALSE_VALID && ? => FALSE_VALID | |
67 | int num_records = 4; | |
68 | auto arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {true, true, true, true}); | |
69 | auto arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
70 | auto exp = MakeArrowArrayBool({false, false, false, false}, {true, true, true, true}); | |
71 | auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
72 | ||
73 | arrow::ArrayVector outputs; | |
74 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
75 | EXPECT_TRUE(status.ok()); | |
76 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
77 | ||
78 | // FALSE_INVALID && ? | |
79 | num_records = 4; | |
80 | arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {false, false, false, false}); | |
81 | arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
82 | exp = MakeArrowArrayBool({false, false, false, false}, {true, false, false, false}); | |
83 | in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
84 | outputs.clear(); | |
85 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
86 | EXPECT_TRUE(status.ok()); | |
87 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
88 | ||
89 | // TRUE_VALID && ? | |
90 | num_records = 4; | |
91 | arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {true, true, true, true}); | |
92 | arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
93 | exp = MakeArrowArrayBool({false, false, true, false}, {true, false, true, false}); | |
94 | in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
95 | outputs.clear(); | |
96 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
97 | EXPECT_TRUE(status.ok()); | |
98 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
99 | ||
100 | // TRUE_INVALID && ? | |
101 | num_records = 4; | |
102 | arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {false, false, false, false}); | |
103 | arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
104 | exp = MakeArrowArrayBool({false, false, false, false}, {true, false, false, false}); | |
105 | in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
106 | outputs.clear(); | |
107 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
108 | EXPECT_TRUE(status.ok()); | |
109 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
110 | } | |
111 | ||
112 | TEST_F(TestBooleanExpr, SimpleOr) { | |
113 | // schema for input fields | |
114 | auto fielda = field("a", int32()); | |
115 | auto fieldb = field("b", int32()); | |
116 | auto schema = arrow::schema({fielda, fieldb}); | |
117 | ||
118 | // output fields | |
119 | auto field_result = field("res", boolean()); | |
120 | ||
121 | // build expression. | |
122 | // (a > 0) || (b > 0) | |
123 | auto node_a = TreeExprBuilder::MakeField(fielda); | |
124 | auto node_b = TreeExprBuilder::MakeField(fieldb); | |
125 | auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0); | |
126 | auto a_gt_0 = | |
127 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean()); | |
128 | auto b_gt_0 = | |
129 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean()); | |
130 | ||
131 | auto node_or = TreeExprBuilder::MakeOr({a_gt_0, b_gt_0}); | |
132 | auto expr = TreeExprBuilder::MakeExpression(node_or, field_result); | |
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()); | |
138 | ||
139 | // TRUE_VALID && ? => TRUE_VALID | |
140 | int num_records = 4; | |
141 | auto arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {true, true, true, true}); | |
142 | auto arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
143 | auto exp = MakeArrowArrayBool({true, true, true, true}, {true, true, true, true}); | |
144 | auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
145 | ||
146 | arrow::ArrayVector outputs; | |
147 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
148 | EXPECT_TRUE(status.ok()); | |
149 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
150 | ||
151 | // TRUE_INVALID && ? | |
152 | num_records = 4; | |
153 | arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {false, false, false, false}); | |
154 | arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
155 | exp = MakeArrowArrayBool({false, false, true, false}, {false, false, true, false}); | |
156 | in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
157 | outputs.clear(); | |
158 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
159 | EXPECT_TRUE(status.ok()); | |
160 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
161 | ||
162 | // FALSE_VALID && ? | |
163 | num_records = 4; | |
164 | arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {true, true, true, true}); | |
165 | arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
166 | exp = MakeArrowArrayBool({false, false, true, false}, {true, false, true, false}); | |
167 | in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
168 | outputs.clear(); | |
169 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
170 | EXPECT_TRUE(status.ok()); | |
171 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
172 | ||
173 | // FALSE_INVALID && ? | |
174 | num_records = 4; | |
175 | arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {false, false, false, false}); | |
176 | arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false}); | |
177 | exp = MakeArrowArrayBool({false, false, true, false}, {false, false, true, false}); | |
178 | in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
179 | outputs.clear(); | |
180 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
181 | EXPECT_TRUE(status.ok()); | |
182 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
183 | } | |
184 | ||
185 | TEST_F(TestBooleanExpr, AndThree) { | |
186 | // schema for input fields | |
187 | auto fielda = field("a", int32()); | |
188 | auto fieldb = field("b", int32()); | |
189 | auto fieldc = field("c", int32()); | |
190 | auto schema = arrow::schema({fielda, fieldb, fieldc}); | |
191 | ||
192 | // output fields | |
193 | auto field_result = field("res", boolean()); | |
194 | ||
195 | // build expression. | |
196 | // (a > 0) && (b > 0) && (c > 0) | |
197 | auto node_a = TreeExprBuilder::MakeField(fielda); | |
198 | auto node_b = TreeExprBuilder::MakeField(fieldb); | |
199 | auto node_c = TreeExprBuilder::MakeField(fieldc); | |
200 | auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0); | |
201 | auto a_gt_0 = | |
202 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean()); | |
203 | auto b_gt_0 = | |
204 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean()); | |
205 | auto c_gt_0 = | |
206 | TreeExprBuilder::MakeFunction("greater_than", {node_c, literal_0}, boolean()); | |
207 | ||
208 | auto node_and = TreeExprBuilder::MakeAnd({a_gt_0, b_gt_0, c_gt_0}); | |
209 | auto expr = TreeExprBuilder::MakeExpression(node_and, field_result); | |
210 | ||
211 | // Build a projector for the expressions. | |
212 | std::shared_ptr<Projector> projector; | |
213 | auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector); | |
214 | EXPECT_TRUE(status.ok()); | |
215 | ||
216 | int num_records = 8; | |
217 | std::vector<bool> validity({true, true, true, true, true, true, true, true}); | |
218 | auto arraya = MakeArrowArrayInt32({2, 2, 2, 0, 2, 0, 0, 0}, validity); | |
219 | auto arrayb = MakeArrowArrayInt32({2, 2, 0, 2, 0, 2, 0, 0}, validity); | |
220 | auto arrayc = MakeArrowArrayInt32({2, 0, 2, 2, 0, 0, 2, 0}, validity); | |
221 | auto exp = MakeArrowArrayBool({true, false, false, false, false, false, false, false}, | |
222 | validity); | |
223 | ||
224 | auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb, arrayc}); | |
225 | ||
226 | arrow::ArrayVector outputs; | |
227 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
228 | EXPECT_TRUE(status.ok()); | |
229 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
230 | } | |
231 | ||
232 | TEST_F(TestBooleanExpr, OrThree) { | |
233 | // schema for input fields | |
234 | auto fielda = field("a", int32()); | |
235 | auto fieldb = field("b", int32()); | |
236 | auto fieldc = field("c", int32()); | |
237 | auto schema = arrow::schema({fielda, fieldb, fieldc}); | |
238 | ||
239 | // output fields | |
240 | auto field_result = field("res", boolean()); | |
241 | ||
242 | // build expression. | |
243 | // (a > 0) || (b > 0) || (c > 0) | |
244 | auto node_a = TreeExprBuilder::MakeField(fielda); | |
245 | auto node_b = TreeExprBuilder::MakeField(fieldb); | |
246 | auto node_c = TreeExprBuilder::MakeField(fieldc); | |
247 | auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0); | |
248 | auto a_gt_0 = | |
249 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean()); | |
250 | auto b_gt_0 = | |
251 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean()); | |
252 | auto c_gt_0 = | |
253 | TreeExprBuilder::MakeFunction("greater_than", {node_c, literal_0}, boolean()); | |
254 | ||
255 | auto node_or = TreeExprBuilder::MakeOr({a_gt_0, b_gt_0, c_gt_0}); | |
256 | auto expr = TreeExprBuilder::MakeExpression(node_or, field_result); | |
257 | ||
258 | // Build a projector for the expressions. | |
259 | std::shared_ptr<Projector> projector; | |
260 | auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector); | |
261 | EXPECT_TRUE(status.ok()); | |
262 | ||
263 | int num_records = 8; | |
264 | std::vector<bool> validity({true, true, true, true, true, true, true, true}); | |
265 | auto arraya = MakeArrowArrayInt32({2, 2, 2, 0, 2, 0, 0, 0}, validity); | |
266 | auto arrayb = MakeArrowArrayInt32({2, 2, 0, 2, 0, 2, 0, 0}, validity); | |
267 | auto arrayc = MakeArrowArrayInt32({2, 0, 2, 2, 0, 0, 2, 0}, validity); | |
268 | auto exp = | |
269 | MakeArrowArrayBool({true, true, true, true, true, true, true, false}, validity); | |
270 | ||
271 | auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb, arrayc}); | |
272 | ||
273 | arrow::ArrayVector outputs; | |
274 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
275 | EXPECT_TRUE(status.ok()); | |
276 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
277 | } | |
278 | ||
279 | TEST_F(TestBooleanExpr, BooleanAndInsideIf) { | |
280 | // schema for input fields | |
281 | auto fielda = field("a", int32()); | |
282 | auto fieldb = field("b", int32()); | |
283 | auto schema = arrow::schema({fielda, fieldb}); | |
284 | ||
285 | // output fields | |
286 | auto field_result = field("res", boolean()); | |
287 | ||
288 | // build expression. | |
289 | // if (a > 2 && b > 2) | |
290 | // a > 3 && b > 3 | |
291 | // else | |
292 | // a > 1 && b > 1 | |
293 | auto node_a = TreeExprBuilder::MakeField(fielda); | |
294 | auto node_b = TreeExprBuilder::MakeField(fieldb); | |
295 | auto literal_1 = TreeExprBuilder::MakeLiteral((int32_t)1); | |
296 | auto literal_2 = TreeExprBuilder::MakeLiteral((int32_t)2); | |
297 | auto literal_3 = TreeExprBuilder::MakeLiteral((int32_t)3); | |
298 | auto a_gt_1 = | |
299 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_1}, boolean()); | |
300 | auto a_gt_2 = | |
301 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_2}, boolean()); | |
302 | auto a_gt_3 = | |
303 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_3}, boolean()); | |
304 | auto b_gt_1 = | |
305 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_1}, boolean()); | |
306 | auto b_gt_2 = | |
307 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_2}, boolean()); | |
308 | auto b_gt_3 = | |
309 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_3}, boolean()); | |
310 | ||
311 | auto and_1 = TreeExprBuilder::MakeAnd({a_gt_1, b_gt_1}); | |
312 | auto and_2 = TreeExprBuilder::MakeAnd({a_gt_2, b_gt_2}); | |
313 | auto and_3 = TreeExprBuilder::MakeAnd({a_gt_3, b_gt_3}); | |
314 | ||
315 | auto node_if = TreeExprBuilder::MakeIf(and_2, and_3, and_1, arrow::boolean()); | |
316 | auto expr = TreeExprBuilder::MakeExpression(node_if, field_result); | |
317 | ||
318 | // Build a projector for the expressions. | |
319 | std::shared_ptr<Projector> projector; | |
320 | auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector); | |
321 | EXPECT_TRUE(status.ok()); | |
322 | ||
323 | int num_records = 4; | |
324 | std::vector<bool> validity({true, true, true, true}); | |
325 | auto arraya = MakeArrowArrayInt32({4, 4, 2, 1}, validity); | |
326 | auto arrayb = MakeArrowArrayInt32({5, 3, 3, 1}, validity); | |
327 | auto exp = MakeArrowArrayBool({true, false, true, false}, validity); | |
328 | ||
329 | auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
330 | ||
331 | arrow::ArrayVector outputs; | |
332 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
333 | EXPECT_TRUE(status.ok()); | |
334 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
335 | } | |
336 | ||
337 | TEST_F(TestBooleanExpr, IfInsideBooleanAnd) { | |
338 | // schema for input fields | |
339 | auto fielda = field("a", int32()); | |
340 | auto fieldb = field("b", int32()); | |
341 | auto schema = arrow::schema({fielda, fieldb}); | |
342 | ||
343 | // output fields | |
344 | auto field_result = field("res", boolean()); | |
345 | ||
346 | // build expression. | |
347 | // (if (a > b) a > 3 else b > 3) && (if (a > b) a > 2 else b > 2) | |
348 | ||
349 | auto node_a = TreeExprBuilder::MakeField(fielda); | |
350 | auto node_b = TreeExprBuilder::MakeField(fieldb); | |
351 | auto literal_2 = TreeExprBuilder::MakeLiteral((int32_t)2); | |
352 | auto literal_3 = TreeExprBuilder::MakeLiteral((int32_t)3); | |
353 | auto a_gt_b = | |
354 | TreeExprBuilder::MakeFunction("greater_than", {node_a, node_b}, boolean()); | |
355 | auto a_gt_2 = | |
356 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_2}, boolean()); | |
357 | auto a_gt_3 = | |
358 | TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_3}, boolean()); | |
359 | auto b_gt_2 = | |
360 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_2}, boolean()); | |
361 | auto b_gt_3 = | |
362 | TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_3}, boolean()); | |
363 | ||
364 | auto if_3 = TreeExprBuilder::MakeIf(a_gt_b, a_gt_3, b_gt_3, arrow::boolean()); | |
365 | auto if_2 = TreeExprBuilder::MakeIf(a_gt_b, a_gt_2, b_gt_2, arrow::boolean()); | |
366 | auto node_and = TreeExprBuilder::MakeAnd({if_3, if_2}); | |
367 | auto expr = TreeExprBuilder::MakeExpression(node_and, field_result); | |
368 | ||
369 | // Build a projector for the expressions. | |
370 | std::shared_ptr<Projector> projector; | |
371 | auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector); | |
372 | EXPECT_TRUE(status.ok()); | |
373 | ||
374 | int num_records = 4; | |
375 | std::vector<bool> validity({true, true, true, true}); | |
376 | auto arraya = MakeArrowArrayInt32({4, 3, 3, 2}, validity); | |
377 | auto arrayb = MakeArrowArrayInt32({3, 4, 2, 3}, validity); | |
378 | auto exp = MakeArrowArrayBool({true, true, false, false}, validity); | |
379 | ||
380 | auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb}); | |
381 | ||
382 | arrow::ArrayVector outputs; | |
383 | status = projector->Evaluate(*in_batch, pool_, &outputs); | |
384 | EXPECT_TRUE(status.ok()); | |
385 | EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0)); | |
386 | } | |
387 | ||
388 | } // namespace gandiva |