arrow/testing/util.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <algorithm>
0021 #include <cstdint>
0022 #include <cstdlib>
0023 #include <cstring>
0024 #include <limits>
0025 #include <memory>
0026 #include <optional>
0027 #include <string>
0028 #include <type_traits>
0029 #include <utility>
0030 #include <vector>
0031
0032 #include "arrow/buffer.h"
0033 #include "arrow/record_batch.h"
0034 #include "arrow/status.h"
0035 #include "arrow/testing/visibility.h"
0036 #include "arrow/type_fwd.h"
0037 #include "arrow/util/macros.h"
0038
0039 namespace arrow {
0040
0041 template <typename T>
0042 Status CopyBufferFromVector(const std::vector<T>& values, MemoryPool* pool,
0043                             std::shared_ptr<Buffer>* result) {
0044   int64_t nbytes = static_cast<int>(values.size()) * sizeof(T);
0045
0046   ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateBuffer(nbytes, pool));
0047   auto immutable_data = reinterpret_cast<const uint8_t*>(values.data());
0048   std::copy(immutable_data, immutable_data + nbytes, buffer->mutable_data());
0049   memset(buffer->mutable_data() + nbytes, 0,
0050          static_cast<size_t>(buffer->capacity() - nbytes));
0051
0052   *result = std::move(buffer);
0053   return Status::OK();
0054 }
0055
0056 // Sets approximately pct_null of the first n bytes in null_bytes to zero
0057 // and the rest to non-zero (true) values.
0058 ARROW_TESTING_EXPORT void random_null_bytes(int64_t n, double pct_null,
0059                                             uint8_t* null_bytes);
0060 ARROW_TESTING_EXPORT void random_is_valid(int64_t n, double pct_null,
0061                                           std::vector<bool>* is_valid,
0062                                           int random_seed = 0);
0063 ARROW_TESTING_EXPORT void random_bytes(int64_t n, uint32_t seed, uint8_t* out);
0064 ARROW_TESTING_EXPORT std::string random_string(int64_t n, uint32_t seed);
0065 ARROW_TESTING_EXPORT int32_t DecimalSize(int32_t precision);
0066 ARROW_TESTING_EXPORT void random_ascii(int64_t n, uint32_t seed, uint8_t* out);
0067 ARROW_TESTING_EXPORT int64_t CountNulls(const std::vector<uint8_t>& valid_bytes);
0068
0069 ARROW_TESTING_EXPORT Status MakeRandomByteBuffer(int64_t length, MemoryPool* pool,
0070                                                  std::shared_ptr<ResizableBuffer>* out,
0071                                                  uint32_t seed = 0);
0072
0073 ARROW_TESTING_EXPORT uint64_t random_seed();
0074
0075 #define DECL_T() typedef typename TestFixture::T T;
0076
0077 #define DECL_TYPE() typedef typename TestFixture::Type Type;
0078
0079 // ----------------------------------------------------------------------
0080 // A RecordBatchReader for serving a sequence of in-memory record batches
0081
0082 class BatchIterator : public RecordBatchReader {
0083  public:
0084   BatchIterator(const std::shared_ptr<Schema>& schema,
0085                 const std::vector<std::shared_ptr<RecordBatch>>& batches)
0086       : schema_(schema), batches_(batches), position_(0) {}
0087
0088   std::shared_ptr<Schema> schema() const override { return schema_; }
0089
0090   Status ReadNext(std::shared_ptr<RecordBatch>* out) override {
0091     if (position_ >= batches_.size()) {
0092       *out = nullptr;
0093     } else {
0094       *out = batches_[position_++];
0095     }
0096     return Status::OK();
0097   }
0098
0099  private:
0100   std::shared_ptr<Schema> schema_;
0101   std::vector<std::shared_ptr<RecordBatch>> batches_;
0102   size_t position_;
0103 };
0104
0105 static inline std::vector<std::shared_ptr<DataType> (*)(FieldVector, std::vector<int8_t>)>
0106 UnionTypeFactories() {
0107   return {sparse_union, dense_union};
0108 }
0109
0110 // Return the value of the ARROW_TEST_DATA environment variable or return error
0111 // Status
0112 ARROW_TESTING_EXPORT Status GetTestResourceRoot(std::string*);
0113
0114 // Return the value of the ARROW_TIMEZONE_DATABASE environment variable
0115 ARROW_TESTING_EXPORT std::optional<std::string> GetTestTimezoneDatabaseRoot();
0116
0117 // Set the Timezone database based on the ARROW_TIMEZONE_DATABASE env variable
0118 // This is only relevant on Windows, since other OSs have compatible databases built-in
0119 ARROW_TESTING_EXPORT Status InitTestTimezoneDatabase();
0120
0121 // Get a TCP port number to listen on.  This is a different number every time,
0122 // as reusing the same port across tests can produce spurious bind errors on
0123 // Windows.
0124 ARROW_TESTING_EXPORT int GetListenPort();
0125
0126 // Get a IPv4 "address:port" to listen on.  The address will be a loopback
0127 // address.  Compared to GetListenPort(), this will minimize the risk of
0128 // port conflicts.
0129 ARROW_TESTING_EXPORT std::string GetListenAddress();
0130
0131 // Get a "host:port" to listen on. Compared to GetListenAddress(), this function would use
0132 // the host passed in.
0133 ARROW_TESTING_EXPORT std::string GetListenAddress(const std::string& host);
0134
0135 ARROW_TESTING_EXPORT
0136 const std::vector<std::shared_ptr<DataType>>& all_dictionary_index_types();
0137
0138 // Get a list of supported hardware flags from the given candidates.
0139 // The result will always contain 0, meaning no optional CPU feature enabled at all.
0140 ARROW_TESTING_EXPORT
0141 std::vector<int64_t> GetSupportedHardwareFlags(
0142     const std::vector<int64_t>& candidate_flags);
0143
0144 }  // namespace arrow