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
9 // http://www.apache.org/licenses/LICENSE-2.0
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
25 #include "arrow/record_batch.h"
26 #include "arrow/testing/gtest_util.h"
27 #include "arrow/testing/visibility.h"
28 #include "arrow/type.h"
32 class ARROW_TESTING_EXPORT ConstantArrayGenerator
{
34 /// \brief Generates a constant BooleanArray
36 /// \param[in] size the size of the array to generate
37 /// \param[in] value to repeat
39 /// \return a generated Array
40 static std::shared_ptr
<arrow::Array
> Boolean(int64_t size
, bool value
= false);
42 /// \brief Generates a constant UInt8Array
44 /// \param[in] size the size of the array to generate
45 /// \param[in] value to repeat
47 /// \return a generated Array
48 static std::shared_ptr
<arrow::Array
> UInt8(int64_t size
, uint8_t value
= 0);
50 /// \brief Generates a constant Int8Array
52 /// \param[in] size the size of the array to generate
53 /// \param[in] value to repeat
55 /// \return a generated Array
56 static std::shared_ptr
<arrow::Array
> Int8(int64_t size
, int8_t value
= 0);
58 /// \brief Generates a constant UInt16Array
60 /// \param[in] size the size of the array to generate
61 /// \param[in] value to repeat
63 /// \return a generated Array
64 static std::shared_ptr
<arrow::Array
> UInt16(int64_t size
, uint16_t value
= 0);
66 /// \brief Generates a constant UInt16Array
68 /// \param[in] size the size of the array to generate
69 /// \param[in] value to repeat
71 /// \return a generated Array
72 static std::shared_ptr
<arrow::Array
> Int16(int64_t size
, int16_t value
= 0);
74 /// \brief Generates a constant UInt32Array
76 /// \param[in] size the size of the array to generate
77 /// \param[in] value to repeat
79 /// \return a generated Array
80 static std::shared_ptr
<arrow::Array
> UInt32(int64_t size
, uint32_t value
= 0);
82 /// \brief Generates a constant UInt32Array
84 /// \param[in] size the size of the array to generate
85 /// \param[in] value to repeat
87 /// \return a generated Array
88 static std::shared_ptr
<arrow::Array
> Int32(int64_t size
, int32_t value
= 0);
90 /// \brief Generates a constant UInt64Array
92 /// \param[in] size the size of the array to generate
93 /// \param[in] value to repeat
95 /// \return a generated Array
96 static std::shared_ptr
<arrow::Array
> UInt64(int64_t size
, uint64_t value
= 0);
98 /// \brief Generates a constant UInt64Array
100 /// \param[in] size the size of the array to generate
101 /// \param[in] value to repeat
103 /// \return a generated Array
104 static std::shared_ptr
<arrow::Array
> Int64(int64_t size
, int64_t value
= 0);
106 /// \brief Generates a constant Float32Array
108 /// \param[in] size the size of the array to generate
109 /// \param[in] value to repeat
111 /// \return a generated Array
112 static std::shared_ptr
<arrow::Array
> Float32(int64_t size
, float value
= 0);
114 /// \brief Generates a constant Float64Array
116 /// \param[in] size the size of the array to generate
117 /// \param[in] value to repeat
119 /// \return a generated Array
120 static std::shared_ptr
<arrow::Array
> Float64(int64_t size
, double value
= 0);
122 /// \brief Generates a constant StringArray
124 /// \param[in] size the size of the array to generate
125 /// \param[in] value to repeat
127 /// \return a generated Array
128 static std::shared_ptr
<arrow::Array
> String(int64_t size
, std::string value
= "");
130 template <typename ArrowType
, typename CType
= typename
ArrowType::c_type
>
131 static std::shared_ptr
<arrow::Array
> Numeric(int64_t size
, CType value
= 0) {
132 switch (ArrowType::type_id
) {
134 return Boolean(size
, static_cast<bool>(value
));
136 return UInt8(size
, static_cast<uint8_t>(value
));
138 return Int8(size
, static_cast<int8_t>(value
));
140 return UInt16(size
, static_cast<uint16_t>(value
));
142 return Int16(size
, static_cast<int16_t>(value
));
144 return UInt32(size
, static_cast<uint32_t>(value
));
146 return Int32(size
, static_cast<int32_t>(value
));
148 return UInt64(size
, static_cast<uint64_t>(value
));
150 return Int64(size
, static_cast<int64_t>(value
));
152 return Float32(size
, static_cast<float>(value
));
154 return Float64(size
, static_cast<double>(value
));
160 /// \brief Generates a constant Array of zeroes
162 /// \param[in] size the size of the array to generate
163 /// \param[in] type the type of the Array
165 /// \return a generated Array
166 static std::shared_ptr
<arrow::Array
> Zeroes(int64_t size
,
167 const std::shared_ptr
<DataType
>& type
) {
168 switch (type
->id()) {
170 return std::make_shared
<NullArray
>(size
);
172 return Boolean(size
);
191 case Type::TIMESTAMP
: {
192 EXPECT_OK_AND_ASSIGN(auto viewed
, Int64(size
)->View(type
));
195 case Type::INTERVAL_DAY_TIME
:
196 case Type::INTERVAL_MONTHS
:
199 EXPECT_OK_AND_ASSIGN(auto viewed
, Int32(size
)->View(type
));
203 return Float32(size
);
205 return Float64(size
);
213 /// \brief Generates a RecordBatch of zeroes
215 /// \param[in] size the size of the array to generate
216 /// \param[in] schema to conform to
218 /// This function is handy to return of RecordBatch of a desired shape.
220 /// \return a generated RecordBatch
221 static std::shared_ptr
<arrow::RecordBatch
> Zeroes(
222 int64_t size
, const std::shared_ptr
<Schema
>& schema
) {
223 std::vector
<std::shared_ptr
<Array
>> arrays
;
225 for (const auto& field
: schema
->fields()) {
226 arrays
.emplace_back(Zeroes(size
, field
->type()));
229 return RecordBatch::Make(schema
, size
, arrays
);
232 /// \brief Generates a RecordBatchReader by repeating a RecordBatch
234 /// \param[in] n_batch the number of times it repeats batch
235 /// \param[in] batch the RecordBatch to repeat
237 /// \return a generated RecordBatchReader
238 static std::shared_ptr
<arrow::RecordBatchReader
> Repeat(
239 int64_t n_batch
, const std::shared_ptr
<RecordBatch
> batch
) {
240 std::vector
<std::shared_ptr
<RecordBatch
>> batches(static_cast<size_t>(n_batch
),
242 return *RecordBatchReader::Make(batches
);
245 /// \brief Generates a RecordBatchReader of zeroes batches
247 /// \param[in] n_batch the number of RecordBatch
248 /// \param[in] batch_size the size of each RecordBatch
249 /// \param[in] schema to conform to
251 /// \return a generated RecordBatchReader
252 static std::shared_ptr
<arrow::RecordBatchReader
> Zeroes(
253 int64_t n_batch
, int64_t batch_size
, const std::shared_ptr
<Schema
>& schema
) {
254 return Repeat(n_batch
, Zeroes(batch_size
, schema
));
259 Result
<std::shared_ptr
<Array
>> ScalarVectorToArray(const ScalarVector
& scalars
);