]>
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 | #pragma once | |
19 | ||
20 | #include <algorithm> | |
21 | #include <cstdint> | |
22 | #include <cstdlib> | |
23 | #include <cstring> | |
24 | #include <functional> | |
25 | #include <memory> | |
26 | #include <string> | |
27 | #include <type_traits> | |
28 | #include <utility> | |
29 | #include <vector> | |
30 | ||
31 | #include <gtest/gtest.h> | |
32 | ||
33 | #include "arrow/array/builder_binary.h" | |
34 | #include "arrow/array/builder_primitive.h" | |
35 | #include "arrow/array/builder_time.h" | |
36 | #include "arrow/result.h" | |
37 | #include "arrow/status.h" | |
38 | #include "arrow/testing/gtest_compat.h" | |
39 | #include "arrow/testing/util.h" | |
40 | #include "arrow/testing/visibility.h" | |
41 | #include "arrow/type_fwd.h" | |
42 | #include "arrow/type_traits.h" | |
43 | #include "arrow/util/bit_util.h" | |
44 | #include "arrow/util/macros.h" | |
45 | #include "arrow/util/string_builder.h" | |
46 | #include "arrow/util/type_fwd.h" | |
47 | ||
48 | // NOTE: failing must be inline in the macros below, to get correct file / line number | |
49 | // reporting on test failures. | |
50 | ||
51 | // NOTE: using a for loop for this macro allows extra failure messages to be | |
52 | // appended with operator<< | |
53 | #define ASSERT_RAISES(ENUM, expr) \ | |
54 | for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); \ | |
55 | !_st.Is##ENUM();) \ | |
56 | FAIL() << "Expected '" ARROW_STRINGIFY(expr) "' to fail with " ARROW_STRINGIFY( \ | |
57 | ENUM) ", but got " \ | |
58 | << _st.ToString() | |
59 | ||
60 | #define ASSERT_RAISES_WITH_MESSAGE(ENUM, message, expr) \ | |
61 | do { \ | |
62 | auto _res = (expr); \ | |
63 | ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \ | |
64 | if (!_st.Is##ENUM()) { \ | |
65 | FAIL() << "Expected '" ARROW_STRINGIFY(expr) "' to fail with " ARROW_STRINGIFY( \ | |
66 | ENUM) ", but got " \ | |
67 | << _st.ToString(); \ | |
68 | } \ | |
69 | ASSERT_EQ((message), _st.ToString()); \ | |
70 | } while (false) | |
71 | ||
72 | #define EXPECT_RAISES_WITH_MESSAGE_THAT(ENUM, matcher, expr) \ | |
73 | do { \ | |
74 | auto _res = (expr); \ | |
75 | ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \ | |
76 | EXPECT_TRUE(_st.Is##ENUM()) << "Expected '" ARROW_STRINGIFY(expr) "' to fail with " \ | |
77 | << ARROW_STRINGIFY(ENUM) ", but got " << _st.ToString(); \ | |
78 | EXPECT_THAT(_st.ToString(), (matcher)); \ | |
79 | } while (false) | |
80 | ||
81 | #define EXPECT_RAISES_WITH_CODE_AND_MESSAGE_THAT(code, matcher, expr) \ | |
82 | do { \ | |
83 | auto _res = (expr); \ | |
84 | ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \ | |
85 | EXPECT_EQ(_st.CodeAsString(), Status::CodeAsString(code)); \ | |
86 | EXPECT_THAT(_st.ToString(), (matcher)); \ | |
87 | } while (false) | |
88 | ||
89 | #define ASSERT_OK(expr) \ | |
90 | for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); !_st.ok();) \ | |
91 | FAIL() << "'" ARROW_STRINGIFY(expr) "' failed with " << _st.ToString() | |
92 | ||
93 | #define ASSERT_OK_NO_THROW(expr) ASSERT_NO_THROW(ASSERT_OK(expr)) | |
94 | ||
95 | #define ARROW_EXPECT_OK(expr) \ | |
96 | do { \ | |
97 | auto _res = (expr); \ | |
98 | ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \ | |
99 | EXPECT_TRUE(_st.ok()) << "'" ARROW_STRINGIFY(expr) "' failed with " \ | |
100 | << _st.ToString(); \ | |
101 | } while (false) | |
102 | ||
103 | #define ASSERT_NOT_OK(expr) \ | |
104 | for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); _st.ok();) \ | |
105 | FAIL() << "'" ARROW_STRINGIFY(expr) "' did not failed" << _st.ToString() | |
106 | ||
107 | #define ABORT_NOT_OK(expr) \ | |
108 | do { \ | |
109 | auto _res = (expr); \ | |
110 | ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \ | |
111 | if (ARROW_PREDICT_FALSE(!_st.ok())) { \ | |
112 | _st.Abort(); \ | |
113 | } \ | |
114 | } while (false); | |
115 | ||
116 | #define ASSIGN_OR_HANDLE_ERROR_IMPL(handle_error, status_name, lhs, rexpr) \ | |
117 | auto&& status_name = (rexpr); \ | |
118 | handle_error(status_name.status()); \ | |
119 | lhs = std::move(status_name).ValueOrDie(); | |
120 | ||
121 | #define ASSERT_OK_AND_ASSIGN(lhs, rexpr) \ | |
122 | ASSIGN_OR_HANDLE_ERROR_IMPL( \ | |
123 | ASSERT_OK, ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), lhs, rexpr); | |
124 | ||
125 | #define ASSIGN_OR_ABORT(lhs, rexpr) \ | |
126 | ASSIGN_OR_HANDLE_ERROR_IMPL(ABORT_NOT_OK, \ | |
127 | ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), \ | |
128 | lhs, rexpr); | |
129 | ||
130 | #define EXPECT_OK_AND_ASSIGN(lhs, rexpr) \ | |
131 | ASSIGN_OR_HANDLE_ERROR_IMPL(ARROW_EXPECT_OK, \ | |
132 | ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), \ | |
133 | lhs, rexpr); | |
134 | ||
135 | #define ASSERT_OK_AND_EQ(expected, expr) \ | |
136 | do { \ | |
137 | ASSERT_OK_AND_ASSIGN(auto _actual, (expr)); \ | |
138 | ASSERT_EQ(expected, _actual); \ | |
139 | } while (0) | |
140 | ||
141 | // A generalized version of GTest's SCOPED_TRACE that takes arbitrary arguments. | |
142 | // ARROW_SCOPED_TRACE("some variable = ", some_variable, ...) | |
143 | ||
144 | #define ARROW_SCOPED_TRACE(...) SCOPED_TRACE(::arrow::util::StringBuilder(__VA_ARGS__)) | |
145 | ||
146 | namespace arrow { | |
147 | ||
148 | // ---------------------------------------------------------------------- | |
149 | // Useful testing::Types declarations | |
150 | ||
151 | inline void PrintTo(StatusCode code, std::ostream* os) { | |
152 | *os << Status::CodeAsString(code); | |
153 | } | |
154 | ||
155 | using NumericArrowTypes = | |
156 | ::testing::Types<UInt8Type, UInt16Type, UInt32Type, UInt64Type, Int8Type, Int16Type, | |
157 | Int32Type, Int64Type, FloatType, DoubleType>; | |
158 | ||
159 | using RealArrowTypes = ::testing::Types<FloatType, DoubleType>; | |
160 | ||
161 | using IntegralArrowTypes = ::testing::Types<UInt8Type, UInt16Type, UInt32Type, UInt64Type, | |
162 | Int8Type, Int16Type, Int32Type, Int64Type>; | |
163 | ||
164 | using PhysicalIntegralArrowTypes = | |
165 | ::testing::Types<UInt8Type, UInt16Type, UInt32Type, UInt64Type, Int8Type, Int16Type, | |
166 | Int32Type, Int64Type, Date32Type, Date64Type, Time32Type, Time64Type, | |
167 | TimestampType, MonthIntervalType>; | |
168 | ||
169 | using PrimitiveArrowTypes = | |
170 | ::testing::Types<BooleanType, Int8Type, UInt8Type, Int16Type, UInt16Type, Int32Type, | |
171 | UInt32Type, Int64Type, UInt64Type, FloatType, DoubleType>; | |
172 | ||
173 | using TemporalArrowTypes = | |
174 | ::testing::Types<Date32Type, Date64Type, TimestampType, Time32Type, Time64Type>; | |
175 | ||
176 | using DecimalArrowTypes = ::testing::Types<Decimal128Type, Decimal256Type>; | |
177 | ||
178 | using BinaryArrowTypes = | |
179 | ::testing::Types<BinaryType, LargeBinaryType, StringType, LargeStringType>; | |
180 | ||
181 | using StringArrowTypes = ::testing::Types<StringType, LargeStringType>; | |
182 | ||
183 | using ListArrowTypes = ::testing::Types<ListType, LargeListType>; | |
184 | ||
185 | using UnionArrowTypes = ::testing::Types<SparseUnionType, DenseUnionType>; | |
186 | ||
187 | class Array; | |
188 | class ChunkedArray; | |
189 | class RecordBatch; | |
190 | class Table; | |
191 | struct Datum; | |
192 | ||
193 | ARROW_TESTING_EXPORT | |
194 | std::vector<Type::type> AllTypeIds(); | |
195 | ||
196 | #define ASSERT_ARRAYS_EQUAL(lhs, rhs) AssertArraysEqual((lhs), (rhs)) | |
197 | #define ASSERT_BATCHES_EQUAL(lhs, rhs) AssertBatchesEqual((lhs), (rhs)) | |
198 | #define ASSERT_BATCHES_APPROX_EQUAL(lhs, rhs) AssertBatchesApproxEqual((lhs), (rhs)) | |
199 | #define ASSERT_TABLES_EQUAL(lhs, rhs) AssertTablesEqual((lhs), (rhs)) | |
200 | ||
201 | // If verbose is true, then the arrays will be pretty printed | |
202 | ARROW_TESTING_EXPORT void AssertArraysEqual(const Array& expected, const Array& actual, | |
203 | bool verbose = false, | |
204 | const EqualOptions& options = {}); | |
205 | ARROW_TESTING_EXPORT void AssertArraysApproxEqual(const Array& expected, | |
206 | const Array& actual, | |
207 | bool verbose = false, | |
208 | const EqualOptions& options = {}); | |
209 | // Returns true when values are both null | |
210 | ARROW_TESTING_EXPORT void AssertScalarsEqual( | |
211 | const Scalar& expected, const Scalar& actual, bool verbose = false, | |
212 | const EqualOptions& options = EqualOptions::Defaults()); | |
213 | ARROW_TESTING_EXPORT void AssertScalarsApproxEqual( | |
214 | const Scalar& expected, const Scalar& actual, bool verbose = false, | |
215 | const EqualOptions& options = EqualOptions::Defaults()); | |
216 | ARROW_TESTING_EXPORT void AssertBatchesEqual(const RecordBatch& expected, | |
217 | const RecordBatch& actual, | |
218 | bool check_metadata = false); | |
219 | ARROW_TESTING_EXPORT void AssertBatchesApproxEqual(const RecordBatch& expected, | |
220 | const RecordBatch& actual); | |
221 | ARROW_TESTING_EXPORT void AssertChunkedEqual(const ChunkedArray& expected, | |
222 | const ChunkedArray& actual); | |
223 | ARROW_TESTING_EXPORT void AssertChunkedEqual(const ChunkedArray& actual, | |
224 | const ArrayVector& expected); | |
225 | // Like ChunkedEqual, but permits different chunk layout | |
226 | ARROW_TESTING_EXPORT void AssertChunkedEquivalent(const ChunkedArray& expected, | |
227 | const ChunkedArray& actual); | |
228 | ARROW_TESTING_EXPORT void AssertChunkedApproxEquivalent( | |
229 | const ChunkedArray& expected, const ChunkedArray& actual, | |
230 | const EqualOptions& equal_options = EqualOptions::Defaults()); | |
231 | ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, | |
232 | const std::vector<uint8_t>& expected); | |
233 | ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, | |
234 | const std::string& expected); | |
235 | ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, const Buffer& expected); | |
236 | ||
237 | ARROW_TESTING_EXPORT void AssertTypeEqual(const DataType& lhs, const DataType& rhs, | |
238 | bool check_metadata = false); | |
239 | ARROW_TESTING_EXPORT void AssertTypeEqual(const std::shared_ptr<DataType>& lhs, | |
240 | const std::shared_ptr<DataType>& rhs, | |
241 | bool check_metadata = false); | |
242 | ARROW_TESTING_EXPORT void AssertFieldEqual(const Field& lhs, const Field& rhs, | |
243 | bool check_metadata = false); | |
244 | ARROW_TESTING_EXPORT void AssertFieldEqual(const std::shared_ptr<Field>& lhs, | |
245 | const std::shared_ptr<Field>& rhs, | |
246 | bool check_metadata = false); | |
247 | ARROW_TESTING_EXPORT void AssertSchemaEqual(const Schema& lhs, const Schema& rhs, | |
248 | bool check_metadata = false); | |
249 | ARROW_TESTING_EXPORT void AssertSchemaEqual(const std::shared_ptr<Schema>& lhs, | |
250 | const std::shared_ptr<Schema>& rhs, | |
251 | bool check_metadata = false); | |
252 | ||
253 | ARROW_TESTING_EXPORT void AssertTypeNotEqual(const DataType& lhs, const DataType& rhs, | |
254 | bool check_metadata = false); | |
255 | ARROW_TESTING_EXPORT void AssertTypeNotEqual(const std::shared_ptr<DataType>& lhs, | |
256 | const std::shared_ptr<DataType>& rhs, | |
257 | bool check_metadata = false); | |
258 | ARROW_TESTING_EXPORT void AssertFieldNotEqual(const Field& lhs, const Field& rhs, | |
259 | bool check_metadata = false); | |
260 | ARROW_TESTING_EXPORT void AssertFieldNotEqual(const std::shared_ptr<Field>& lhs, | |
261 | const std::shared_ptr<Field>& rhs, | |
262 | bool check_metadata = false); | |
263 | ARROW_TESTING_EXPORT void AssertSchemaNotEqual(const Schema& lhs, const Schema& rhs, | |
264 | bool check_metadata = false); | |
265 | ARROW_TESTING_EXPORT void AssertSchemaNotEqual(const std::shared_ptr<Schema>& lhs, | |
266 | const std::shared_ptr<Schema>& rhs, | |
267 | bool check_metadata = false); | |
268 | ||
269 | ARROW_TESTING_EXPORT Result<util::optional<std::string>> PrintArrayDiff( | |
270 | const ChunkedArray& expected, const ChunkedArray& actual); | |
271 | ||
272 | ARROW_TESTING_EXPORT void AssertTablesEqual(const Table& expected, const Table& actual, | |
273 | bool same_chunk_layout = true, | |
274 | bool flatten = false); | |
275 | ||
276 | ARROW_TESTING_EXPORT void AssertDatumsEqual(const Datum& expected, const Datum& actual, | |
277 | bool verbose = false); | |
278 | ARROW_TESTING_EXPORT void AssertDatumsApproxEqual( | |
279 | const Datum& expected, const Datum& actual, bool verbose = false, | |
280 | const EqualOptions& options = EqualOptions::Defaults()); | |
281 | ||
282 | template <typename C_TYPE> | |
283 | void AssertNumericDataEqual(const C_TYPE* raw_data, | |
284 | const std::vector<C_TYPE>& expected_values) { | |
285 | for (auto expected : expected_values) { | |
286 | ASSERT_EQ(expected, *raw_data); | |
287 | ++raw_data; | |
288 | } | |
289 | } | |
290 | ||
291 | ARROW_TESTING_EXPORT void CompareBatch(const RecordBatch& left, const RecordBatch& right, | |
292 | bool compare_metadata = true); | |
293 | ||
294 | ARROW_TESTING_EXPORT void ApproxCompareBatch(const RecordBatch& left, | |
295 | const RecordBatch& right, | |
296 | bool compare_metadata = true); | |
297 | ||
298 | // Check if the padding of the buffers of the array is zero. | |
299 | // Also cause valgrind warnings if the padding bytes are uninitialized. | |
300 | ARROW_TESTING_EXPORT void AssertZeroPadded(const Array& array); | |
301 | ||
302 | // Check if the valid buffer bytes are initialized | |
303 | // and cause valgrind warnings otherwise. | |
304 | ARROW_TESTING_EXPORT void TestInitialized(const ArrayData& array); | |
305 | ARROW_TESTING_EXPORT void TestInitialized(const Array& array); | |
306 | ||
307 | template <typename BuilderType> | |
308 | void FinishAndCheckPadding(BuilderType* builder, std::shared_ptr<Array>* out) { | |
309 | ASSERT_OK_AND_ASSIGN(*out, builder->Finish()); | |
310 | AssertZeroPadded(**out); | |
311 | TestInitialized(**out); | |
312 | } | |
313 | ||
314 | #define DECL_T() typedef typename TestFixture::T T; | |
315 | ||
316 | #define DECL_TYPE() typedef typename TestFixture::Type Type; | |
317 | ||
318 | // ArrayFromJSON: construct an Array from a simple JSON representation | |
319 | ||
320 | ARROW_TESTING_EXPORT | |
321 | std::shared_ptr<Array> ArrayFromJSON(const std::shared_ptr<DataType>&, | |
322 | util::string_view json); | |
323 | ||
324 | ARROW_TESTING_EXPORT | |
325 | std::shared_ptr<Array> DictArrayFromJSON(const std::shared_ptr<DataType>& type, | |
326 | util::string_view indices_json, | |
327 | util::string_view dictionary_json); | |
328 | ||
329 | ARROW_TESTING_EXPORT | |
330 | std::shared_ptr<RecordBatch> RecordBatchFromJSON(const std::shared_ptr<Schema>&, | |
331 | util::string_view); | |
332 | ||
333 | ARROW_TESTING_EXPORT | |
334 | std::shared_ptr<ChunkedArray> ChunkedArrayFromJSON(const std::shared_ptr<DataType>&, | |
335 | const std::vector<std::string>& json); | |
336 | ||
337 | ARROW_TESTING_EXPORT | |
338 | std::shared_ptr<Scalar> ScalarFromJSON(const std::shared_ptr<DataType>&, | |
339 | util::string_view json); | |
340 | ||
341 | ARROW_TESTING_EXPORT | |
342 | std::shared_ptr<Scalar> DictScalarFromJSON(const std::shared_ptr<DataType>&, | |
343 | util::string_view index_json, | |
344 | util::string_view dictionary_json); | |
345 | ||
346 | ARROW_TESTING_EXPORT | |
347 | std::shared_ptr<Table> TableFromJSON(const std::shared_ptr<Schema>&, | |
348 | const std::vector<std::string>& json); | |
349 | ||
350 | // ArrayFromVector: construct an Array from vectors of C values | |
351 | ||
352 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
353 | void ArrayFromVector(const std::shared_ptr<DataType>& type, | |
354 | const std::vector<bool>& is_valid, const std::vector<C_TYPE>& values, | |
355 | std::shared_ptr<Array>* out) { | |
356 | auto type_id = TYPE::type_id; | |
357 | ASSERT_EQ(type_id, type->id()) | |
358 | << "template parameter and concrete DataType instance don't agree"; | |
359 | ||
360 | std::unique_ptr<ArrayBuilder> builder_ptr; | |
361 | ASSERT_OK(MakeBuilder(default_memory_pool(), type, &builder_ptr)); | |
362 | // Get the concrete builder class to access its Append() specializations | |
363 | auto& builder = dynamic_cast<typename TypeTraits<TYPE>::BuilderType&>(*builder_ptr); | |
364 | ||
365 | for (size_t i = 0; i < values.size(); ++i) { | |
366 | if (is_valid[i]) { | |
367 | ASSERT_OK(builder.Append(values[i])); | |
368 | } else { | |
369 | ASSERT_OK(builder.AppendNull()); | |
370 | } | |
371 | } | |
372 | ASSERT_OK(builder.Finish(out)); | |
373 | } | |
374 | ||
375 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
376 | void ArrayFromVector(const std::shared_ptr<DataType>& type, | |
377 | const std::vector<C_TYPE>& values, std::shared_ptr<Array>* out) { | |
378 | auto type_id = TYPE::type_id; | |
379 | ASSERT_EQ(type_id, type->id()) | |
380 | << "template parameter and concrete DataType instance don't agree"; | |
381 | ||
382 | std::unique_ptr<ArrayBuilder> builder_ptr; | |
383 | ASSERT_OK(MakeBuilder(default_memory_pool(), type, &builder_ptr)); | |
384 | // Get the concrete builder class to access its Append() specializations | |
385 | auto& builder = dynamic_cast<typename TypeTraits<TYPE>::BuilderType&>(*builder_ptr); | |
386 | ||
387 | for (size_t i = 0; i < values.size(); ++i) { | |
388 | ASSERT_OK(builder.Append(values[i])); | |
389 | } | |
390 | ASSERT_OK(builder.Finish(out)); | |
391 | } | |
392 | ||
393 | // Overloads without a DataType argument, for parameterless types | |
394 | ||
395 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
396 | void ArrayFromVector(const std::vector<bool>& is_valid, const std::vector<C_TYPE>& values, | |
397 | std::shared_ptr<Array>* out) { | |
398 | auto type = TypeTraits<TYPE>::type_singleton(); | |
399 | ArrayFromVector<TYPE, C_TYPE>(type, is_valid, values, out); | |
400 | } | |
401 | ||
402 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
403 | void ArrayFromVector(const std::vector<C_TYPE>& values, std::shared_ptr<Array>* out) { | |
404 | auto type = TypeTraits<TYPE>::type_singleton(); | |
405 | ArrayFromVector<TYPE, C_TYPE>(type, values, out); | |
406 | } | |
407 | ||
408 | // ChunkedArrayFromVector: construct a ChunkedArray from vectors of C values | |
409 | ||
410 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
411 | void ChunkedArrayFromVector(const std::shared_ptr<DataType>& type, | |
412 | const std::vector<std::vector<bool>>& is_valid, | |
413 | const std::vector<std::vector<C_TYPE>>& values, | |
414 | std::shared_ptr<ChunkedArray>* out) { | |
415 | ArrayVector chunks; | |
416 | ASSERT_EQ(is_valid.size(), values.size()); | |
417 | for (size_t i = 0; i < values.size(); ++i) { | |
418 | std::shared_ptr<Array> array; | |
419 | ArrayFromVector<TYPE, C_TYPE>(type, is_valid[i], values[i], &array); | |
420 | chunks.push_back(array); | |
421 | } | |
422 | *out = std::make_shared<ChunkedArray>(chunks); | |
423 | } | |
424 | ||
425 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
426 | void ChunkedArrayFromVector(const std::shared_ptr<DataType>& type, | |
427 | const std::vector<std::vector<C_TYPE>>& values, | |
428 | std::shared_ptr<ChunkedArray>* out) { | |
429 | ArrayVector chunks; | |
430 | for (size_t i = 0; i < values.size(); ++i) { | |
431 | std::shared_ptr<Array> array; | |
432 | ArrayFromVector<TYPE, C_TYPE>(type, values[i], &array); | |
433 | chunks.push_back(array); | |
434 | } | |
435 | *out = std::make_shared<ChunkedArray>(chunks); | |
436 | } | |
437 | ||
438 | // Overloads without a DataType argument, for parameterless types | |
439 | ||
440 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
441 | void ChunkedArrayFromVector(const std::vector<std::vector<bool>>& is_valid, | |
442 | const std::vector<std::vector<C_TYPE>>& values, | |
443 | std::shared_ptr<ChunkedArray>* out) { | |
444 | auto type = TypeTraits<TYPE>::type_singleton(); | |
445 | ChunkedArrayFromVector<TYPE, C_TYPE>(type, is_valid, values, out); | |
446 | } | |
447 | ||
448 | template <typename TYPE, typename C_TYPE = typename TYPE::c_type> | |
449 | void ChunkedArrayFromVector(const std::vector<std::vector<C_TYPE>>& values, | |
450 | std::shared_ptr<ChunkedArray>* out) { | |
451 | auto type = TypeTraits<TYPE>::type_singleton(); | |
452 | ChunkedArrayFromVector<TYPE, C_TYPE>(type, values, out); | |
453 | } | |
454 | ||
455 | template <typename T> | |
456 | static inline Status GetBitmapFromVector(const std::vector<T>& is_valid, | |
457 | std::shared_ptr<Buffer>* result) { | |
458 | size_t length = is_valid.size(); | |
459 | ||
460 | ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateEmptyBitmap(length)); | |
461 | ||
462 | uint8_t* bitmap = buffer->mutable_data(); | |
463 | for (size_t i = 0; i < static_cast<size_t>(length); ++i) { | |
464 | if (is_valid[i]) { | |
465 | BitUtil::SetBit(bitmap, i); | |
466 | } | |
467 | } | |
468 | ||
469 | *result = buffer; | |
470 | return Status::OK(); | |
471 | } | |
472 | ||
473 | template <typename T> | |
474 | inline void BitmapFromVector(const std::vector<T>& is_valid, | |
475 | std::shared_ptr<Buffer>* out) { | |
476 | ASSERT_OK(GetBitmapFromVector(is_valid, out)); | |
477 | } | |
478 | ||
479 | // Given an array, return a new identical array except for one validity bit | |
480 | // set to a new value. | |
481 | // This is useful to force the underlying "value" of null entries to otherwise | |
482 | // invalid data and check that errors don't get reported. | |
483 | ARROW_TESTING_EXPORT | |
484 | std::shared_ptr<Array> TweakValidityBit(const std::shared_ptr<Array>& array, | |
485 | int64_t index, bool validity); | |
486 | ||
487 | ARROW_TESTING_EXPORT | |
488 | void SleepFor(double seconds); | |
489 | ||
490 | // Sleeps for a very small amount of time. The thread will be yielded | |
491 | // at least once ensuring that context switches could happen. It is intended | |
492 | // to be used for stress testing parallel code and shouldn't be assumed to do any | |
493 | // reliable timing. | |
494 | ARROW_TESTING_EXPORT | |
495 | void SleepABit(); | |
496 | ||
497 | // Wait until predicate is true or timeout in seconds expires. | |
498 | ARROW_TESTING_EXPORT | |
499 | void BusyWait(double seconds, std::function<bool()> predicate); | |
500 | ||
501 | ARROW_TESTING_EXPORT | |
502 | Future<> SleepAsync(double seconds); | |
503 | ||
504 | // \see SleepABit | |
505 | ARROW_TESTING_EXPORT | |
506 | Future<> SleepABitAsync(); | |
507 | ||
508 | template <typename T> | |
509 | std::vector<T> IteratorToVector(Iterator<T> iterator) { | |
510 | EXPECT_OK_AND_ASSIGN(auto out, iterator.ToVector()); | |
511 | return out; | |
512 | } | |
513 | ||
514 | ARROW_TESTING_EXPORT | |
515 | bool LocaleExists(const char* locale); | |
516 | ||
517 | // A RAII-style object that switches to a new locale, and switches back | |
518 | // to the old locale when going out of scope. Doesn't do anything if the | |
519 | // new locale doesn't exist on the local machine. | |
520 | // ATTENTION: may crash with an assertion failure on Windows debug builds. | |
521 | // See ARROW-6108, also https://gerrit.libreoffice.org/#/c/54110/ | |
522 | class ARROW_TESTING_EXPORT LocaleGuard { | |
523 | public: | |
524 | explicit LocaleGuard(const char* new_locale); | |
525 | ~LocaleGuard(); | |
526 | ||
527 | protected: | |
528 | class Impl; | |
529 | std::unique_ptr<Impl> impl_; | |
530 | }; | |
531 | ||
532 | class ARROW_TESTING_EXPORT EnvVarGuard { | |
533 | public: | |
534 | EnvVarGuard(const std::string& name, const std::string& value); | |
535 | ~EnvVarGuard(); | |
536 | ||
537 | protected: | |
538 | const std::string name_; | |
539 | std::string old_value_; | |
540 | bool was_set_; | |
541 | }; | |
542 | ||
543 | namespace internal { | |
544 | class SignalHandler; | |
545 | } | |
546 | ||
547 | class ARROW_TESTING_EXPORT SignalHandlerGuard { | |
548 | public: | |
549 | typedef void (*Callback)(int); | |
550 | ||
551 | SignalHandlerGuard(int signum, Callback cb); | |
552 | SignalHandlerGuard(int signum, const internal::SignalHandler& handler); | |
553 | ~SignalHandlerGuard(); | |
554 | ||
555 | protected: | |
556 | struct Impl; | |
557 | std::unique_ptr<Impl> impl_; | |
558 | }; | |
559 | ||
560 | #ifndef ARROW_LARGE_MEMORY_TESTS | |
561 | #define LARGE_MEMORY_TEST(name) DISABLED_##name | |
562 | #else | |
563 | #define LARGE_MEMORY_TEST(name) name | |
564 | #endif | |
565 | ||
566 | inline void PrintTo(const Status& st, std::ostream* os) { *os << st.ToString(); } | |
567 | ||
568 | template <typename T> | |
569 | void PrintTo(const Result<T>& result, std::ostream* os) { | |
570 | if (result.ok()) { | |
571 | ::testing::internal::UniversalPrint(result.ValueOrDie(), os); | |
572 | } else { | |
573 | *os << result.status(); | |
574 | } | |
575 | } | |
576 | ||
577 | // A data type with only move constructors (no copy, no default). | |
578 | struct MoveOnlyDataType { | |
579 | explicit MoveOnlyDataType(int x) : data(new int(x)) {} | |
580 | ||
581 | MoveOnlyDataType(const MoveOnlyDataType& other) = delete; | |
582 | MoveOnlyDataType& operator=(const MoveOnlyDataType& other) = delete; | |
583 | ||
584 | MoveOnlyDataType(MoveOnlyDataType&& other) { MoveFrom(&other); } | |
585 | MoveOnlyDataType& operator=(MoveOnlyDataType&& other) { | |
586 | MoveFrom(&other); | |
587 | return *this; | |
588 | } | |
589 | ||
590 | MoveOnlyDataType& operator=(int x) { | |
591 | if (data != nullptr) { | |
592 | delete data; | |
593 | } | |
594 | data = new int(x); | |
595 | return *this; | |
596 | } | |
597 | ||
598 | ~MoveOnlyDataType() { Destroy(); } | |
599 | ||
600 | void Destroy() { | |
601 | if (data != nullptr) { | |
602 | delete data; | |
603 | data = nullptr; | |
604 | moves = -1; | |
605 | } | |
606 | } | |
607 | ||
608 | void MoveFrom(MoveOnlyDataType* other) { | |
609 | Destroy(); | |
610 | data = other->data; | |
611 | other->data = nullptr; | |
612 | moves = other->moves + 1; | |
613 | } | |
614 | ||
615 | int ToInt() const { return data == nullptr ? -42 : *data; } | |
616 | ||
617 | bool operator==(const MoveOnlyDataType& other) const { | |
618 | return data != nullptr && other.data != nullptr && *data == *other.data; | |
619 | } | |
620 | bool operator<(const MoveOnlyDataType& other) const { | |
621 | return data == nullptr || (other.data != nullptr && *data < *other.data); | |
622 | } | |
623 | ||
624 | bool operator==(int other) const { return data != nullptr && *data == other; } | |
625 | friend bool operator==(int left, const MoveOnlyDataType& right) { | |
626 | return right == left; | |
627 | } | |
628 | ||
629 | int* data = nullptr; | |
630 | int moves = 0; | |
631 | }; | |
632 | ||
633 | // A task that blocks until unlocked. Useful for timing tests. | |
634 | class ARROW_TESTING_EXPORT GatingTask { | |
635 | public: | |
636 | explicit GatingTask(double timeout_seconds = 10); | |
637 | /// \brief During destruction we wait for all pending tasks to finish | |
638 | ~GatingTask(); | |
639 | ||
640 | /// \brief Creates a new waiting task (presumably to spawn on a thread). It will return | |
641 | /// invalid if the timeout arrived before the unlock. The task will not complete until | |
642 | /// unlocked or timed out | |
643 | /// | |
644 | /// Note: The GatingTask must outlive any Task instances | |
645 | std::function<void()> Task(); | |
646 | /// \brief Creates a new waiting task as a future. The future will not complete | |
647 | /// until unlocked. | |
648 | Future<> AsyncTask(); | |
649 | /// \brief Waits until at least count tasks are running. | |
650 | Status WaitForRunning(int count); | |
651 | /// \brief Unlocks all waiting tasks. Returns an invalid status if any waiting task has | |
652 | /// timed out | |
653 | Status Unlock(); | |
654 | ||
655 | static std::shared_ptr<GatingTask> Make(double timeout_seconds = 10); | |
656 | ||
657 | private: | |
658 | class Impl; | |
659 | std::shared_ptr<Impl> impl_; | |
660 | }; | |
661 | ||
662 | } // namespace arrow | |
663 | ||
664 | namespace nonstd { | |
665 | namespace sv_lite { | |
666 | ||
667 | // Without this hint, GTest will print string_views as a container of char | |
668 | template <class Char, class Traits = std::char_traits<Char>> | |
669 | void PrintTo(const basic_string_view<Char, Traits>& view, std::ostream* os) { | |
670 | *os << view; | |
671 | } | |
672 | ||
673 | } // namespace sv_lite | |
674 | ||
675 | namespace optional_lite { | |
676 | ||
677 | template <typename T> | |
678 | void PrintTo(const optional<T>& opt, std::ostream* os) { | |
679 | if (opt.has_value()) { | |
680 | *os << "{"; | |
681 | ::testing::internal::UniversalPrint(*opt, os); | |
682 | *os << "}"; | |
683 | } else { | |
684 | *os << "nullopt"; | |
685 | } | |
686 | } | |
687 | ||
688 | inline void PrintTo(const decltype(nullopt)&, std::ostream* os) { *os << "nullopt"; } | |
689 | ||
690 | } // namespace optional_lite | |
691 | } // namespace nonstd |