| 44 | } |
| 45 | |
| 46 | TEST_F(TestAnnotator, TestAdd) { |
| 47 | Annotator annotator; |
| 48 | |
| 49 | auto field_a = arrow::field("a", arrow::int32()); |
| 50 | auto field_b = arrow::field("b", arrow::int32()); |
| 51 | auto in_schema = arrow::schema({field_a, field_b}); |
| 52 | auto field_sum = arrow::field("sum", arrow::int32()); |
| 53 | |
| 54 | FieldDescriptorPtr desc_a = annotator.CheckAndAddInputFieldDescriptor(field_a); |
| 55 | EXPECT_EQ(desc_a->field(), field_a); |
| 56 | EXPECT_EQ(desc_a->data_idx(), 0); |
| 57 | EXPECT_EQ(desc_a->validity_idx(), 1); |
| 58 | |
| 59 | // duplicate add shouldn't cause a new descriptor. |
| 60 | FieldDescriptorPtr dup = annotator.CheckAndAddInputFieldDescriptor(field_a); |
| 61 | EXPECT_EQ(dup, desc_a); |
| 62 | EXPECT_EQ(dup->validity_idx(), desc_a->validity_idx()); |
| 63 | |
| 64 | FieldDescriptorPtr desc_b = annotator.CheckAndAddInputFieldDescriptor(field_b); |
| 65 | EXPECT_EQ(desc_b->field(), field_b); |
| 66 | EXPECT_EQ(desc_b->data_idx(), 2); |
| 67 | EXPECT_EQ(desc_b->validity_idx(), 3); |
| 68 | |
| 69 | FieldDescriptorPtr desc_sum = annotator.AddOutputFieldDescriptor(field_sum); |
| 70 | EXPECT_EQ(desc_sum->field(), field_sum); |
| 71 | EXPECT_EQ(desc_sum->data_idx(), 4); |
| 72 | EXPECT_EQ(desc_sum->validity_idx(), 5); |
| 73 | EXPECT_EQ(desc_sum->data_buffer_ptr_idx(), 6); |
| 74 | |
| 75 | // prepare record batch |
| 76 | int num_records = 100; |
| 77 | auto arrow_v0 = MakeInt32Array(num_records); |
| 78 | auto arrow_v1 = MakeInt32Array(num_records); |
| 79 | |
| 80 | // prepare input record batch |
| 81 | auto record_batch = |
| 82 | arrow::RecordBatch::Make(in_schema, num_records, {arrow_v0, arrow_v1}); |
| 83 | |
| 84 | auto arrow_sum = MakeInt32Array(num_records); |
| 85 | EvalBatchPtr batch = annotator.PrepareEvalBatch(*record_batch, {arrow_sum->data()}); |
| 86 | EXPECT_EQ(batch->GetNumBuffers(), 7); |
| 87 | |
| 88 | auto buffers = batch->GetBufferArray(); |
| 89 | EXPECT_EQ(buffers[desc_a->validity_idx()], arrow_v0->data()->buffers.at(0)->data()); |
| 90 | EXPECT_EQ(buffers[desc_a->data_idx()], arrow_v0->data()->buffers.at(1)->data()); |
| 91 | EXPECT_EQ(buffers[desc_b->validity_idx()], arrow_v1->data()->buffers.at(0)->data()); |
| 92 | EXPECT_EQ(buffers[desc_b->data_idx()], arrow_v1->data()->buffers.at(1)->data()); |
| 93 | EXPECT_EQ(buffers[desc_sum->validity_idx()], arrow_sum->data()->buffers.at(0)->data()); |
| 94 | EXPECT_EQ(buffers[desc_sum->data_idx()], arrow_sum->data()->buffers.at(1)->data()); |
| 95 | EXPECT_EQ(buffers[desc_sum->data_buffer_ptr_idx()], |
| 96 | reinterpret_cast<uint8_t*>(arrow_sum->data()->buffers.at(1).get())); |
| 97 | |
| 98 | auto bitmaps = batch->GetLocalBitMapArray(); |
| 99 | EXPECT_EQ(bitmaps, nullptr); |
| 100 | } |
| 101 | |
| 102 | } // namespace gandiva |
nothing calls this directly
no test coverage detected