/usr/local/lib64/python3.6/site-packages/pyarrow/include/arrow/testing
Edit: /usr/local/lib64/python3.6/site-packages/pyarrow/include/arrow/testing/generator.h (8645B)
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#pragma once
#include
#include
#include
#include
#include "arrow/record_batch.h"
#include "arrow/testing/gtest_util.h"
#include "arrow/testing/visibility.h"
#include "arrow/type.h"
namespace arrow {
class ARROW_TESTING_EXPORT ConstantArrayGenerator {
public:
/// \brief Generates a constant BooleanArray
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr Boolean(int64_t size, bool value = false);
/// \brief Generates a constant UInt8Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr UInt8(int64_t size, uint8_t value = 0);
/// \brief Generates a constant Int8Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr Int8(int64_t size, int8_t value = 0);
/// \brief Generates a constant UInt16Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr UInt16(int64_t size, uint16_t value = 0);
/// \brief Generates a constant UInt16Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr Int16(int64_t size, int16_t value = 0);
/// \brief Generates a constant UInt32Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr UInt32(int64_t size, uint32_t value = 0);
/// \brief Generates a constant UInt32Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr Int32(int64_t size, int32_t value = 0);
/// \brief Generates a constant UInt64Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr UInt64(int64_t size, uint64_t value = 0);
/// \brief Generates a constant UInt64Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr Int64(int64_t size, int64_t value = 0);
/// \brief Generates a constant Float32Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr Float32(int64_t size, float value = 0);
/// \brief Generates a constant Float64Array
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr Float64(int64_t size, double value = 0);
/// \brief Generates a constant StringArray
///
/// \param[in] size the size of the array to generate
/// \param[in] value to repeat
///
/// \return a generated Array
static std::shared_ptr String(int64_t size, std::string value = "");
template
static std::shared_ptr Numeric(int64_t size, CType value = 0) {
switch (ArrowType::type_id) {
case Type::BOOL:
return Boolean(size, static_cast(value));
case Type::UINT8:
return UInt8(size, static_cast(value));
case Type::INT8:
return Int8(size, static_cast(value));
case Type::UINT16:
return UInt16(size, static_cast(value));
case Type::INT16:
return Int16(size, static_cast(value));
case Type::UINT32:
return UInt32(size, static_cast(value));
case Type::INT32:
return Int32(size, static_cast(value));
case Type::UINT64:
return UInt64(size, static_cast(value));
case Type::INT64:
return Int64(size, static_cast(value));
case Type::FLOAT:
return Float32(size, static_cast(value));
case Type::DOUBLE:
return Float64(size, static_cast(value));
default:
return nullptr;
}
}
/// \brief Generates a constant Array of zeroes
///
/// \param[in] size the size of the array to generate
/// \param[in] type the type of the Array
///
/// \return a generated Array
static std::shared_ptr Zeroes(int64_t size,
const std::shared_ptr& type) {
switch (type->id()) {
case Type::NA:
return std::make_shared(size);
case Type::BOOL:
return Boolean(size);
case Type::UINT8:
return UInt8(size);
case Type::INT8:
return Int8(size);
case Type::UINT16:
return UInt16(size);
case Type::INT16:
return Int16(size);
case Type::UINT32:
return UInt32(size);
case Type::INT32:
return Int32(size);
case Type::UINT64:
return UInt64(size);
case Type::INT64:
return Int64(size);
case Type::TIME64:
case Type::DATE64:
case Type::TIMESTAMP: {
EXPECT_OK_AND_ASSIGN(auto viewed, Int64(size)->View(type));
return viewed;
}
case Type::INTERVAL_DAY_TIME:
case Type::INTERVAL_MONTHS:
case Type::TIME32:
case Type::DATE32: {
EXPECT_OK_AND_ASSIGN(auto viewed, Int32(size)->View(type));
return viewed;
}
case Type::FLOAT:
return Float32(size);
case Type::DOUBLE:
return Float64(size);
case Type::STRING:
return String(size);
default:
return nullptr;
}
}
/// \brief Generates a RecordBatch of zeroes
///
/// \param[in] size the size of the array to generate
/// \param[in] schema to conform to
///
/// This function is handy to return of RecordBatch of a desired shape.
///
/// \return a generated RecordBatch
static std::shared_ptr Zeroes(
int64_t size, const std::shared_ptr& schema) {
std::vector> arrays;
for (const auto& field : schema->fields()) {
arrays.emplace_back(Zeroes(size, field->type()));
}
return RecordBatch::Make(schema, size, arrays);
}
/// \brief Generates a RecordBatchReader by repeating a RecordBatch
///
/// \param[in] n_batch the number of times it repeats batch
/// \param[in] batch the RecordBatch to repeat
///
/// \return a generated RecordBatchReader
static std::shared_ptr Repeat(
int64_t n_batch, const std::shared_ptr batch) {
std::vector> batches(static_cast(n_batch),
batch);
return *RecordBatchReader::Make(batches);
}
/// \brief Generates a RecordBatchReader of zeroes batches
///
/// \param[in] n_batch the number of RecordBatch
/// \param[in] batch_size the size of each RecordBatch
/// \param[in] schema to conform to
///
/// \return a generated RecordBatchReader
static std::shared_ptr Zeroes(
int64_t n_batch, int64_t batch_size, const std::shared_ptr& schema) {
return Repeat(n_batch, Zeroes(batch_size, schema));
}
};
ARROW_TESTING_EXPORT
Result> ScalarVectorToArray(const ScalarVector& scalars);
} // namespace arrow