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 <cstdint>
21 #include <memory>
22 #include <random>
23 #include <string>
24 #include <utility>
25 #include <vector>
26
27 #include <gtest/gtest.h>
28
29 #include "arrow/array.h"
30 #include "arrow/array/builder_binary.h"
31 #include "arrow/array/builder_primitive.h"
32 #include "arrow/testing/gtest_util.h"
33 #include "arrow/testing/util.h"
34 #include "arrow/type_fwd.h"
35
36 namespace arrow {
37
38 class TestBase : public ::testing::Test {
39 public:
SetUp()40 void SetUp() {
41 pool_ = default_memory_pool();
42 random_seed_ = 0;
43 }
44
MakeRandomNullBitmap(int64_t length,int64_t null_count)45 std::shared_ptr<Buffer> MakeRandomNullBitmap(int64_t length, int64_t null_count) {
46 const int64_t null_nbytes = BitUtil::BytesForBits(length);
47
48 std::shared_ptr<Buffer> null_bitmap = *AllocateBuffer(null_nbytes, pool_);
49 memset(null_bitmap->mutable_data(), 255, null_nbytes);
50 for (int64_t i = 0; i < null_count; i++) {
51 BitUtil::ClearBit(null_bitmap->mutable_data(), i * (length / null_count));
52 }
53 return null_bitmap;
54 }
55
56 template <typename ArrayType>
57 inline std::shared_ptr<Array> MakeRandomArray(int64_t length, int64_t null_count = 0);
58
59 protected:
60 uint32_t random_seed_;
61 MemoryPool* pool_;
62 };
63
64 template <typename ArrayType>
MakeRandomArray(int64_t length,int64_t null_count)65 std::shared_ptr<Array> TestBase::MakeRandomArray(int64_t length, int64_t null_count) {
66 const int64_t data_nbytes = length * sizeof(typename ArrayType::value_type);
67 auto data = *AllocateBuffer(data_nbytes, pool_);
68
69 // Fill with random data
70 random_bytes(data_nbytes, random_seed_++, data->mutable_data());
71 std::shared_ptr<Buffer> null_bitmap = MakeRandomNullBitmap(length, null_count);
72
73 return std::make_shared<ArrayType>(length, std::move(data), null_bitmap, null_count);
74 }
75
76 template <>
77 inline std::shared_ptr<Array> TestBase::MakeRandomArray<NullArray>(int64_t length,
78 int64_t null_count) {
79 return std::make_shared<NullArray>(length);
80 }
81
82 template <>
83 inline std::shared_ptr<Array> TestBase::MakeRandomArray<FixedSizeBinaryArray>(
84 int64_t length, int64_t null_count) {
85 const int byte_width = 10;
86 std::shared_ptr<Buffer> null_bitmap = MakeRandomNullBitmap(length, null_count);
87 auto data = *AllocateBuffer(byte_width * length, pool_);
88
89 ::arrow::random_bytes(data->size(), 0, data->mutable_data());
90 return std::make_shared<FixedSizeBinaryArray>(fixed_size_binary(byte_width), length,
91 std::move(data), null_bitmap, null_count);
92 }
93
94 template <>
95 inline std::shared_ptr<Array> TestBase::MakeRandomArray<BinaryArray>(int64_t length,
96 int64_t null_count) {
97 std::vector<uint8_t> valid_bytes(length, 1);
98 for (int64_t i = 0; i < null_count; i++) {
99 valid_bytes[i * 2] = 0;
100 }
101 BinaryBuilder builder(pool_);
102
103 const int kBufferSize = 10;
104 uint8_t buffer[kBufferSize];
105 for (int64_t i = 0; i < length; i++) {
106 if (!valid_bytes[i]) {
107 ARROW_EXPECT_OK(builder.AppendNull());
108 } else {
109 ::arrow::random_bytes(kBufferSize, static_cast<uint32_t>(i), buffer);
110 ARROW_EXPECT_OK(builder.Append(buffer, kBufferSize));
111 }
112 }
113
114 std::shared_ptr<Array> out;
115 ARROW_EXPECT_OK(builder.Finish(&out));
116 return out;
117 }
118
119 class TestBuilder : public ::testing::Test {
120 public:
SetUp()121 void SetUp() { pool_ = default_memory_pool(); }
122
123 protected:
124 MemoryPool* pool_;
125 std::shared_ptr<DataType> type_;
126 };
127
128 } // namespace arrow
129