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