File indexing completed on 2025-08-28 08:27:01
0001
0002
0003
0004
0005
0006
0007
0008
0009
0010
0011
0012
0013
0014
0015
0016
0017
0018 #pragma once
0019
0020 #include <algorithm>
0021 #include <cstdint>
0022 #include <cstdlib>
0023 #include <cstring>
0024 #include <functional>
0025 #include <memory>
0026 #include <optional>
0027 #include <string>
0028 #include <string_view>
0029 #include <type_traits>
0030 #include <utility>
0031 #include <vector>
0032
0033 #include <gtest/gtest.h>
0034
0035 #include "arrow/compare.h"
0036 #include "arrow/result.h"
0037 #include "arrow/status.h"
0038 #include "arrow/testing/gtest_compat.h"
0039 #include "arrow/testing/visibility.h"
0040 #include "arrow/type_fwd.h"
0041 #include "arrow/type_traits.h"
0042 #include "arrow/util/macros.h"
0043 #include "arrow/util/string_builder.h"
0044 #include "arrow/util/type_fwd.h"
0045
0046
0047
0048
0049
0050
0051 #define ASSERT_RAISES(ENUM, expr) \
0052 for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); \
0053 !_st.Is##ENUM();) \
0054 FAIL() << "Expected '" ARROW_STRINGIFY(expr) "' to fail with " ARROW_STRINGIFY( \
0055 ENUM) ", but got " \
0056 << _st.ToString()
0057
0058 #define ASSERT_RAISES_WITH_MESSAGE(ENUM, message, expr) \
0059 do { \
0060 auto _res = (expr); \
0061 ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \
0062 if (!_st.Is##ENUM()) { \
0063 FAIL() << "Expected '" ARROW_STRINGIFY(expr) "' to fail with " ARROW_STRINGIFY( \
0064 ENUM) ", but got " \
0065 << _st.ToString(); \
0066 } \
0067 ASSERT_EQ((message), _st.ToStringWithoutContextLines()); \
0068 } while (false)
0069
0070 #define EXPECT_RAISES_WITH_MESSAGE_THAT(ENUM, matcher, expr) \
0071 do { \
0072 auto _res = (expr); \
0073 ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \
0074 EXPECT_TRUE(_st.Is##ENUM()) << "Expected '" ARROW_STRINGIFY(expr) "' to fail with " \
0075 << ARROW_STRINGIFY(ENUM) ", but got " << _st.ToString(); \
0076 EXPECT_THAT(_st.ToStringWithoutContextLines(), (matcher)); \
0077 } while (false)
0078
0079 #define EXPECT_RAISES_WITH_CODE_AND_MESSAGE_THAT(code, matcher, expr) \
0080 do { \
0081 auto _res = (expr); \
0082 ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \
0083 EXPECT_EQ(_st.CodeAsString(), Status::CodeAsString(code)); \
0084 EXPECT_THAT(_st.ToStringWithoutContextLines(), (matcher)); \
0085 } while (false)
0086
0087 #define ASSERT_OK(expr) \
0088 for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); !_st.ok();) \
0089 FAIL() << "'" ARROW_STRINGIFY(expr) "' failed with " << _st.ToString()
0090
0091 #define ASSERT_OK_NO_THROW(expr) ASSERT_NO_THROW(ASSERT_OK(expr))
0092
0093 #define ARROW_EXPECT_OK(expr) \
0094 do { \
0095 auto _res = (expr); \
0096 ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \
0097 EXPECT_TRUE(_st.ok()) << "'" ARROW_STRINGIFY(expr) "' failed with " \
0098 << _st.ToString(); \
0099 } while (false)
0100
0101 #define ASSERT_NOT_OK(expr) \
0102 for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); _st.ok();) \
0103 FAIL() << "'" ARROW_STRINGIFY(expr) "' did not failed" << _st.ToString()
0104
0105 #define ABORT_NOT_OK(expr) \
0106 do { \
0107 auto _res = (expr); \
0108 ::arrow::Status _st = ::arrow::internal::GenericToStatus(_res); \
0109 if (ARROW_PREDICT_FALSE(!_st.ok())) { \
0110 _st.Abort(); \
0111 } \
0112 } while (false);
0113
0114 #define ASSIGN_OR_HANDLE_ERROR_IMPL(handle_error, status_name, lhs, rexpr) \
0115 auto&& status_name = (rexpr); \
0116 handle_error(status_name.status()); \
0117 lhs = std::move(status_name).ValueOrDie();
0118
0119 #define ASSERT_OK_AND_ASSIGN(lhs, rexpr) \
0120 ASSIGN_OR_HANDLE_ERROR_IMPL( \
0121 ASSERT_OK, ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), lhs, rexpr);
0122
0123 #define ASSIGN_OR_ABORT(lhs, rexpr) \
0124 ASSIGN_OR_HANDLE_ERROR_IMPL(ABORT_NOT_OK, \
0125 ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), \
0126 lhs, rexpr);
0127
0128 #define EXPECT_OK_AND_ASSIGN(lhs, rexpr) \
0129 ASSIGN_OR_HANDLE_ERROR_IMPL(ARROW_EXPECT_OK, \
0130 ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), \
0131 lhs, rexpr);
0132
0133 #define ASSERT_OK_AND_EQ(expected, expr) \
0134 do { \
0135 ASSERT_OK_AND_ASSIGN(auto _actual, (expr)); \
0136 ASSERT_EQ(expected, _actual); \
0137 } while (0)
0138
0139
0140
0141
0142 #define ARROW_SCOPED_TRACE(...) SCOPED_TRACE(::arrow::util::StringBuilder(__VA_ARGS__))
0143
0144 namespace arrow {
0145
0146
0147
0148
0149 inline void PrintTo(StatusCode code, std::ostream* os) {
0150 *os << Status::CodeAsString(code);
0151 }
0152
0153 using NumericArrowTypes =
0154 ::testing::Types<UInt8Type, UInt16Type, UInt32Type, UInt64Type, Int8Type, Int16Type,
0155 Int32Type, Int64Type, FloatType, DoubleType>;
0156
0157 using RealArrowTypes = ::testing::Types<FloatType, DoubleType>;
0158
0159 using IntegralArrowTypes = ::testing::Types<UInt8Type, UInt16Type, UInt32Type, UInt64Type,
0160 Int8Type, Int16Type, Int32Type, Int64Type>;
0161
0162 using PhysicalIntegralArrowTypes =
0163 ::testing::Types<UInt8Type, UInt16Type, UInt32Type, UInt64Type, Int8Type, Int16Type,
0164 Int32Type, Int64Type, Date32Type, Date64Type, Time32Type, Time64Type,
0165 TimestampType, MonthIntervalType>;
0166
0167 using PrimitiveArrowTypes =
0168 ::testing::Types<BooleanType, Int8Type, UInt8Type, Int16Type, UInt16Type, Int32Type,
0169 UInt32Type, Int64Type, UInt64Type, FloatType, DoubleType>;
0170
0171 using TemporalArrowTypes =
0172 ::testing::Types<Date32Type, Date64Type, TimestampType, Time32Type, Time64Type>;
0173
0174
0175
0176 using DecimalArrowTypes =
0177 ::testing::Types< Decimal128Type, Decimal256Type>;
0178
0179 using BaseBinaryArrowTypes =
0180 ::testing::Types<BinaryType, LargeBinaryType, StringType, LargeStringType>;
0181
0182 using BaseBinaryOrBinaryViewLikeArrowTypes =
0183 ::testing::Types<BinaryType, LargeBinaryType, BinaryViewType, StringType,
0184 LargeStringType, StringViewType>;
0185
0186 using BinaryArrowTypes = ::testing::Types<BinaryType, LargeBinaryType>;
0187
0188 using StringArrowTypes = ::testing::Types<StringType, LargeStringType>;
0189
0190 using StringOrStringViewArrowTypes =
0191 ::testing::Types<StringType, LargeStringType, StringViewType>;
0192
0193 using ListArrowTypes = ::testing::Types<ListType, LargeListType>;
0194
0195 using UnionArrowTypes = ::testing::Types<SparseUnionType, DenseUnionType>;
0196
0197 class Array;
0198 class ChunkedArray;
0199 class RecordBatch;
0200 class Table;
0201 struct Datum;
0202
0203 #define ASSERT_ARRAYS_EQUAL(lhs, rhs) AssertArraysEqual((lhs), (rhs))
0204 #define ASSERT_BATCHES_EQUAL(lhs, rhs) AssertBatchesEqual((lhs), (rhs))
0205 #define ASSERT_BATCHES_APPROX_EQUAL(lhs, rhs) AssertBatchesApproxEqual((lhs), (rhs))
0206 #define ASSERT_TABLES_EQUAL(lhs, rhs) AssertTablesEqual((lhs), (rhs))
0207
0208
0209 static inline EqualOptions TestingEqualOptions() {
0210 return EqualOptions{}.nans_equal(true).signed_zeros_equal(false);
0211 }
0212
0213
0214 ARROW_TESTING_EXPORT void AssertArraysEqual(
0215 const Array& expected, const Array& actual, bool verbose = false,
0216 const EqualOptions& options = TestingEqualOptions());
0217 ARROW_TESTING_EXPORT void AssertArraysApproxEqual(
0218 const Array& expected, const Array& actual, bool verbose = false,
0219 const EqualOptions& options = TestingEqualOptions());
0220
0221 ARROW_TESTING_EXPORT void AssertScalarsEqual(
0222 const Scalar& expected, const Scalar& actual, bool verbose = false,
0223 const EqualOptions& options = TestingEqualOptions());
0224 ARROW_TESTING_EXPORT void AssertScalarsApproxEqual(
0225 const Scalar& expected, const Scalar& actual, bool verbose = false,
0226 const EqualOptions& options = TestingEqualOptions());
0227 ARROW_TESTING_EXPORT void AssertBatchesEqual(
0228 const RecordBatch& expected, const RecordBatch& actual, bool check_metadata = false,
0229 const EqualOptions& options = TestingEqualOptions());
0230 ARROW_TESTING_EXPORT void AssertBatchesApproxEqual(
0231 const RecordBatch& expected, const RecordBatch& actual,
0232 const EqualOptions& options = TestingEqualOptions());
0233 ARROW_TESTING_EXPORT void AssertChunkedEqual(
0234 const ChunkedArray& expected, const ChunkedArray& actual,
0235 const EqualOptions& options = TestingEqualOptions());
0236 ARROW_TESTING_EXPORT void AssertChunkedEqual(
0237 const ChunkedArray& actual, const ArrayVector& expected,
0238 const EqualOptions& options = TestingEqualOptions());
0239
0240 ARROW_TESTING_EXPORT void AssertChunkedEquivalent(
0241 const ChunkedArray& expected, const ChunkedArray& actual,
0242 const EqualOptions& options = TestingEqualOptions());
0243 ARROW_TESTING_EXPORT void AssertChunkedApproxEquivalent(
0244 const ChunkedArray& expected, const ChunkedArray& actual,
0245 const EqualOptions& options = TestingEqualOptions());
0246 ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer,
0247 const std::vector<uint8_t>& expected);
0248 ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer,
0249 std::string_view expected);
0250 ARROW_TESTING_EXPORT void AssertBufferEqual(const Buffer& buffer, const Buffer& expected);
0251
0252 ARROW_TESTING_EXPORT void AssertTypeEqual(const DataType& lhs, const DataType& rhs,
0253 bool check_metadata = false);
0254 ARROW_TESTING_EXPORT void AssertTypeEqual(const std::shared_ptr<DataType>& lhs,
0255 const std::shared_ptr<DataType>& rhs,
0256 bool check_metadata = false);
0257 ARROW_TESTING_EXPORT void AssertFieldEqual(const Field& lhs, const Field& rhs,
0258 bool check_metadata = false);
0259 ARROW_TESTING_EXPORT void AssertFieldEqual(const std::shared_ptr<Field>& lhs,
0260 const std::shared_ptr<Field>& rhs,
0261 bool check_metadata = false);
0262 ARROW_TESTING_EXPORT void AssertSchemaEqual(const Schema& lhs, const Schema& rhs,
0263 bool check_metadata = false);
0264 ARROW_TESTING_EXPORT void AssertSchemaEqual(const std::shared_ptr<Schema>& lhs,
0265 const std::shared_ptr<Schema>& rhs,
0266 bool check_metadata = false);
0267
0268 ARROW_TESTING_EXPORT void AssertTypeNotEqual(const DataType& lhs, const DataType& rhs,
0269 bool check_metadata = false);
0270 ARROW_TESTING_EXPORT void AssertTypeNotEqual(const std::shared_ptr<DataType>& lhs,
0271 const std::shared_ptr<DataType>& rhs,
0272 bool check_metadata = false);
0273 ARROW_TESTING_EXPORT void AssertFieldNotEqual(const Field& lhs, const Field& rhs,
0274 bool check_metadata = false);
0275 ARROW_TESTING_EXPORT void AssertFieldNotEqual(const std::shared_ptr<Field>& lhs,
0276 const std::shared_ptr<Field>& rhs,
0277 bool check_metadata = false);
0278 ARROW_TESTING_EXPORT void AssertSchemaNotEqual(const Schema& lhs, const Schema& rhs,
0279 bool check_metadata = false);
0280 ARROW_TESTING_EXPORT void AssertSchemaNotEqual(const std::shared_ptr<Schema>& lhs,
0281 const std::shared_ptr<Schema>& rhs,
0282 bool check_metadata = false);
0283
0284 ARROW_TESTING_EXPORT Result<std::optional<std::string>> PrintArrayDiff(
0285 const ChunkedArray& expected, const ChunkedArray& actual);
0286
0287 ARROW_TESTING_EXPORT void AssertTablesEqual(
0288 const Table& expected, const Table& actual, bool same_chunk_layout = true,
0289 bool flatten = false, const EqualOptions& options = TestingEqualOptions());
0290
0291 ARROW_TESTING_EXPORT void AssertDatumsEqual(
0292 const Datum& expected, const Datum& actual, bool verbose = false,
0293 const EqualOptions& options = TestingEqualOptions());
0294 ARROW_TESTING_EXPORT void AssertDatumsApproxEqual(
0295 const Datum& expected, const Datum& actual, bool verbose = false,
0296 const EqualOptions& options = TestingEqualOptions());
0297
0298 template <typename C_TYPE>
0299 void AssertNumericDataEqual(const C_TYPE* raw_data,
0300 const std::vector<C_TYPE>& expected_values) {
0301 for (auto expected : expected_values) {
0302 ASSERT_EQ(expected, *raw_data);
0303 ++raw_data;
0304 }
0305 }
0306
0307 ARROW_TESTING_EXPORT void CompareBatch(
0308 const RecordBatch& left, const RecordBatch& right, bool compare_metadata = true,
0309 const EqualOptions& options = TestingEqualOptions());
0310
0311 ARROW_TESTING_EXPORT void ApproxCompareBatch(
0312 const RecordBatch& left, const RecordBatch& right, bool compare_metadata = true,
0313 const EqualOptions& options = TestingEqualOptions());
0314
0315
0316
0317 ARROW_TESTING_EXPORT void AssertZeroPadded(const Array& array);
0318
0319
0320
0321 ARROW_TESTING_EXPORT void TestInitialized(const ArrayData& array);
0322 ARROW_TESTING_EXPORT void TestInitialized(const Array& array);
0323
0324 #define DECL_T() typedef typename TestFixture::T T;
0325
0326 #define DECL_TYPE() typedef typename TestFixture::Type Type;
0327
0328
0329
0330 ARROW_TESTING_EXPORT
0331 std::shared_ptr<Array> ArrayFromJSON(const std::shared_ptr<DataType>&,
0332 std::string_view json);
0333
0334 ARROW_TESTING_EXPORT
0335 std::shared_ptr<Array> DictArrayFromJSON(const std::shared_ptr<DataType>& type,
0336 std::string_view indices_json,
0337 std::string_view dictionary_json);
0338
0339 ARROW_TESTING_EXPORT
0340 std::shared_ptr<RecordBatch> RecordBatchFromJSON(const std::shared_ptr<Schema>&,
0341 std::string_view);
0342
0343 ARROW_TESTING_EXPORT
0344 std::shared_ptr<ChunkedArray> ChunkedArrayFromJSON(const std::shared_ptr<DataType>&,
0345 const std::vector<std::string>& json);
0346
0347 ARROW_TESTING_EXPORT
0348 std::shared_ptr<Scalar> ScalarFromJSON(const std::shared_ptr<DataType>&,
0349 std::string_view json);
0350
0351 ARROW_TESTING_EXPORT
0352 std::shared_ptr<Scalar> DictScalarFromJSON(const std::shared_ptr<DataType>&,
0353 std::string_view index_json,
0354 std::string_view dictionary_json);
0355
0356 ARROW_TESTING_EXPORT
0357 std::shared_ptr<Table> TableFromJSON(const std::shared_ptr<Schema>&,
0358 const std::vector<std::string>& json);
0359
0360 ARROW_TESTING_EXPORT
0361 std::shared_ptr<Tensor> TensorFromJSON(const std::shared_ptr<DataType>& type,
0362 std::string_view data, std::string_view shape,
0363 std::string_view strides = "[]",
0364 std::string_view dim_names = "[]");
0365
0366 ARROW_TESTING_EXPORT
0367 std::shared_ptr<Tensor> TensorFromJSON(const std::shared_ptr<DataType>& type,
0368 std::string_view data,
0369 const std::vector<int64_t>& shape,
0370 const std::vector<int64_t>& strides = {},
0371 const std::vector<std::string>& dim_names = {});
0372
0373 ARROW_TESTING_EXPORT
0374 Result<std::shared_ptr<Table>> RunEndEncodeTableColumns(
0375 const Table& table, const std::vector<int>& column_indices);
0376
0377
0378
0379
0380
0381 ARROW_TESTING_EXPORT
0382 std::shared_ptr<Array> TweakValidityBit(const std::shared_ptr<Array>& array,
0383 int64_t index, bool validity);
0384
0385 ARROW_TESTING_EXPORT
0386 void SleepFor(double seconds);
0387
0388
0389
0390
0391
0392 ARROW_TESTING_EXPORT
0393 void SleepABit();
0394
0395
0396 ARROW_TESTING_EXPORT
0397 void BusyWait(double seconds, std::function<bool()> predicate);
0398
0399
0400 ARROW_TESTING_EXPORT
0401 Future<> SleepABitAsync();
0402
0403 ARROW_TESTING_EXPORT bool FileIsClosed(int fd);
0404
0405 template <typename T>
0406 std::vector<T> IteratorToVector(Iterator<T> iterator) {
0407 EXPECT_OK_AND_ASSIGN(auto out, iterator.ToVector());
0408 return out;
0409 }
0410
0411 ARROW_TESTING_EXPORT
0412 bool LocaleExists(const char* locale);
0413
0414 #ifndef _WIN32
0415 ARROW_TESTING_EXPORT
0416 void AssertChildExit(int child_pid, int expected_exit_status = 0);
0417 #endif
0418
0419
0420
0421
0422
0423
0424 class ARROW_TESTING_EXPORT LocaleGuard {
0425 public:
0426 explicit LocaleGuard(const char* new_locale);
0427 ~LocaleGuard();
0428
0429 protected:
0430 class Impl;
0431 std::unique_ptr<Impl> impl_;
0432 };
0433
0434 class ARROW_TESTING_EXPORT EnvVarGuard {
0435 public:
0436 EnvVarGuard(const std::string& name, const std::string& value);
0437 ~EnvVarGuard();
0438
0439 protected:
0440 const std::string name_;
0441 std::string old_value_;
0442 bool was_set_;
0443 };
0444
0445 namespace internal {
0446 class SignalHandler;
0447 }
0448
0449 class ARROW_TESTING_EXPORT SignalHandlerGuard {
0450 public:
0451 typedef void (*Callback)(int);
0452
0453 SignalHandlerGuard(int signum, Callback cb);
0454 SignalHandlerGuard(int signum, const internal::SignalHandler& handler);
0455 ~SignalHandlerGuard();
0456
0457 protected:
0458 struct Impl;
0459 std::unique_ptr<Impl> impl_;
0460 };
0461
0462 #ifndef ARROW_LARGE_MEMORY_TESTS
0463 # define LARGE_MEMORY_TEST(name) DISABLED_##name
0464 #else
0465 # define LARGE_MEMORY_TEST(name) name
0466 #endif
0467
0468 inline void PrintTo(const Status& st, std::ostream* os) { *os << st.ToString(); }
0469
0470 template <typename T>
0471 void PrintTo(const Result<T>& result, std::ostream* os) {
0472 if (result.ok()) {
0473 ::testing::internal::UniversalPrint(result.ValueOrDie(), os);
0474 } else {
0475 *os << result.status();
0476 }
0477 }
0478
0479
0480 struct MoveOnlyDataType {
0481 explicit MoveOnlyDataType(int x) : data(new int(x)) {}
0482
0483 MoveOnlyDataType(const MoveOnlyDataType& other) = delete;
0484 MoveOnlyDataType& operator=(const MoveOnlyDataType& other) = delete;
0485
0486 MoveOnlyDataType(MoveOnlyDataType&& other) { MoveFrom(&other); }
0487 MoveOnlyDataType& operator=(MoveOnlyDataType&& other) {
0488 MoveFrom(&other);
0489 return *this;
0490 }
0491
0492 MoveOnlyDataType& operator=(int x) {
0493 if (data != nullptr) {
0494 delete data;
0495 }
0496 data = new int(x);
0497 return *this;
0498 }
0499
0500 ~MoveOnlyDataType() { Destroy(); }
0501
0502 void Destroy() {
0503 if (data != nullptr) {
0504 delete data;
0505 data = nullptr;
0506 moves = -1;
0507 }
0508 }
0509
0510 void MoveFrom(MoveOnlyDataType* other) {
0511 Destroy();
0512 data = other->data;
0513 other->data = nullptr;
0514 moves = other->moves + 1;
0515 }
0516
0517 int ToInt() const { return data == nullptr ? -42 : *data; }
0518
0519 bool operator==(const MoveOnlyDataType& other) const {
0520 return data != nullptr && other.data != nullptr && *data == *other.data;
0521 }
0522 bool operator<(const MoveOnlyDataType& other) const {
0523 return data == nullptr || (other.data != nullptr && *data < *other.data);
0524 }
0525
0526 bool operator==(int other) const { return data != nullptr && *data == other; }
0527 friend bool operator==(int left, const MoveOnlyDataType& right) {
0528 return right == left;
0529 }
0530
0531 int* data = nullptr;
0532 int moves = 0;
0533 };
0534
0535
0536 class ARROW_TESTING_EXPORT GatingTask {
0537 public:
0538 explicit GatingTask(double timeout_seconds = 10);
0539
0540 ~GatingTask();
0541
0542
0543
0544
0545
0546
0547 std::function<void()> Task();
0548
0549
0550 Future<> AsyncTask();
0551
0552 Status WaitForRunning(int count);
0553
0554
0555 Status Unlock();
0556
0557 static std::shared_ptr<GatingTask> Make(double timeout_seconds = 10);
0558
0559 private:
0560 class Impl;
0561 std::shared_ptr<Impl> impl_;
0562 };
0563
0564
0565
0566
0567 ARROW_TESTING_EXPORT std::shared_ptr<ArrayData> UnalignBuffers(const ArrayData& array);
0568
0569
0570
0571 ARROW_TESTING_EXPORT std::shared_ptr<Array> UnalignBuffers(const Array& array);
0572
0573 }