mirrors/arrow
0
0 Fork 0
mirror of https://github.com/apache/arrow.git synced 2026-10-28 16:47:19 +08:00
Code Activity

Make tests templated

Browse Source
This commit is contained in:
Uwe L. Korn committed 2016-06-03 22:29:36 +02:00
1 parent be6415c393
commit 9b06e417f7
1 file changed
+117 -115
Download Patch File Download Diff File Expand all files Collapse all files
cpp/src/arrow/parquet/parquet-io-test.cc
+117 -115
View File
@@ -18,6 +18,7 @@
#include "gtest/gtest.h"
#include "arrow/test-util.h"
#include "arrow/parquet/test-util.h"
#include "arrow/parquet/reader.h"
#include "arrow/parquet/writer.h"
#include "arrow/types/primitive.h"
@@ -44,48 +45,36 @@ namespace arrow {
namespace parquet {
template <typename ArrowType>
std::shared_ptr<PrimitiveArray> NonNullArray(
size_t size, typename ArrowType::c_type value) {
std::vector<typename ArrowType::c_type> values(size, value);
NumericBuilder<ArrowType> builder(default_memory_pool(), std::make_shared<ArrowType>());
builder.Append(values.data(), values.size());
return std::static_pointer_cast<PrimitiveArray>(builder.Finish());
}
const int SMALL_SIZE = 100;
const int LARGE_SIZE = 10000;
// This helper function only supports (size/2) nulls yet.
template <typename ArrowType>
std::shared_ptr<PrimitiveArray> NullableArray(
size_t size, typename ArrowType::c_type value, size_t num_nulls) {
std::vector<typename ArrowType::c_type> values(size, value);
std::vector<uint8_t> valid_bytes(size, 1);
template <typename TestType>
struct test_traits {};
for (size_t i = 0; i < num_nulls; i++) {
valid_bytes[i * 2] = 0;
}
template <>
struct test_traits<Int32Type> {
static constexpr ParquetType::type parquet_enum = ParquetType::INT32;
};
NumericBuilder<ArrowType> builder(default_memory_pool(), std::make_shared<ArrowType>());
builder.Append(values.data(), values.size(), valid_bytes.data());
return std::static_pointer_cast<PrimitiveArray>(builder.Finish());
}
template <>
struct test_traits<Int64Type> {
static constexpr ParquetType::type parquet_enum = ParquetType::INT64;
};
std::shared_ptr<Column> MakeColumn(const std::string& name,
const std::shared_ptr<PrimitiveArray>& array, bool nullable) {
auto field = std::make_shared<Field>(name, array->type(), nullable);
return std::make_shared<Column>(field, array);
}
template <>
struct test_traits<FloatType> {
static constexpr ParquetType::type parquet_enum = ParquetType::FLOAT;
};
std::shared_ptr<Table> MakeSimpleTable(
const std::shared_ptr<PrimitiveArray>& values, bool nullable) {
std::shared_ptr<Column> column = MakeColumn("col", values, nullable);
std::vector<std::shared_ptr<Column>> columns({column});
std::vector<std::shared_ptr<Field>> fields({column->field()});
auto schema = std::make_shared<Schema>(fields);
return std::make_shared<Table>("table", schema, columns);
}
template <>
struct test_traits<DoubleType> {
static constexpr ParquetType::type parquet_enum = ParquetType::DOUBLE;
};
template <typename TestType>
class TestParquetIO : public ::testing::Test {
public:
typedef typename TestType::c_type T;
virtual void SetUp() {}
std::shared_ptr<GroupNode> MakeSchema(
@@ -113,7 +102,7 @@ class TestParquetIO : public ::testing::Test {
std::unique_ptr<arrow::parquet::FlatColumnReader> column_reader;
ASSERT_NO_THROW(ASSERT_OK(reader.GetFlatColumn(0, &column_reader)));
ASSERT_NE(nullptr, column_reader.get());
ASSERT_OK(column_reader->NextBatch(100, out));
ASSERT_OK(column_reader->NextBatch(SMALL_SIZE, out));
ASSERT_NE(nullptr, out->get());
}
@@ -143,93 +132,104 @@ class TestParquetIO : public ::testing::Test {
return ReaderFromSink();
}
std::unique_ptr<ParquetFileReader> TestFile(std::vector<T>& values, int num_chunks) {
std::shared_ptr<GroupNode> schema =
MakeSchema(test_traits<TestType>::parquet_enum, Repetition::REQUIRED);
std::unique_ptr<ParquetFileWriter> file_writer = MakeWriter(schema);
size_t chunk_size = values.size() / num_chunks;
for (int i = 0; i < num_chunks; i++) {
auto row_group_writer = file_writer->AppendRowGroup(chunk_size);
auto column_writer = static_cast<::parquet::TypedColumnWriter<
::parquet::DataType<test_traits<TestType>::parquet_enum>>*>(
row_group_writer->NextColumn());
T* data = values.data() + i * chunk_size;
column_writer->WriteBatch(chunk_size, nullptr, nullptr, data);
column_writer->Close();
row_group_writer->Close();
}
file_writer->Close();
return ReaderFromSink();
}
std::shared_ptr<InMemoryOutputStream> sink_;
};
TEST_F(TestParquetIO, SingleColumnInt64Read) {
std::vector<int64_t> values(100, 128);
std::unique_ptr<ParquetFileReader> file_reader = Int64File(values, 1);
typedef ::testing::Types<Int32Type, Int64Type, FloatType, DoubleType> TestTypes;
TYPED_TEST_CASE(TestParquetIO, TestTypes);
TYPED_TEST(TestParquetIO, SingleColumnRequiredRead) {
std::vector<typename TypeParam::c_type> values(SMALL_SIZE, 128);
std::unique_ptr<ParquetFileReader> file_reader = this->TestFile(values, 1);
std::shared_ptr<Array> out;
ReadSingleColumnFile(std::move(file_reader), &out);
this->ReadSingleColumnFile(std::move(file_reader), &out);
Int64Array* out_array = static_cast<Int64Array*>(out.get());
for (size_t i = 0; i < values.size(); i++) {
EXPECT_EQ(values[i], out_array->raw_data()[i]);
}
ExpectArray<typename TypeParam::c_type>(values.data(), out.get());
}
TEST_F(TestParquetIO, SingleColumnInt64TableRead) {
std::vector<int64_t> values(100, 128);
std::unique_ptr<ParquetFileReader> file_reader = Int64File(values, 1);
TYPED_TEST(TestParquetIO, SingleColumnRequiredTableRead) {
std::vector<typename TypeParam::c_type> values(SMALL_SIZE, 128);
std::unique_ptr<ParquetFileReader> file_reader = this->TestFile(values, 1);
std::shared_ptr<Table> out;
ReadTableFromFile(std::move(file_reader), &out);
this->ReadTableFromFile(std::move(file_reader), &out);
ASSERT_EQ(1, out->num_columns());
ASSERT_EQ(100, out->num_rows());
ASSERT_EQ(SMALL_SIZE, out->num_rows());
std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data();
ASSERT_EQ(1, chunked_array->num_chunks());
Int64Array* out_array = static_cast<Int64Array*>(chunked_array->chunk(0).get());
for (size_t i = 0; i < values.size(); i++) {
EXPECT_EQ(values[i], out_array->raw_data()[i]);
}
ExpectArray<typename TypeParam::c_type>(values.data(), chunked_array->chunk(0).get());
}
TEST_F(TestParquetIO, SingleColumnInt64ChunkedRead) {
std::vector<int64_t> values(100, 128);
std::unique_ptr<ParquetFileReader> file_reader = Int64File(values, 4);
TYPED_TEST(TestParquetIO, SingleColumnRequiredChunkedRead) {
std::vector<typename TypeParam::c_type> values(SMALL_SIZE, 128);
std::unique_ptr<ParquetFileReader> file_reader = this->TestFile(values, 4);
std::shared_ptr<Array> out;
ReadSingleColumnFile(std::move(file_reader), &out);
this->ReadSingleColumnFile(std::move(file_reader), &out);
Int64Array* out_array = static_cast<Int64Array*>(out.get());
for (size_t i = 0; i < values.size(); i++) {
EXPECT_EQ(values[i], out_array->raw_data()[i]);
}
ExpectArray<typename TypeParam::c_type>(values.data(), out.get());
}
TEST_F(TestParquetIO, SingleColumnInt64ChunkedTableRead) {
std::vector<int64_t> values(100, 128);
std::unique_ptr<ParquetFileReader> file_reader = Int64File(values, 4);
TYPED_TEST(TestParquetIO, SingleColumnRequiredChunkedTableRead) {
std::vector<typename TypeParam::c_type> values(SMALL_SIZE, 128);
std::unique_ptr<ParquetFileReader> file_reader = this->TestFile(values, 4);
std::shared_ptr<Table> out;
ReadTableFromFile(std::move(file_reader), &out);
this->ReadTableFromFile(std::move(file_reader), &out);
ASSERT_EQ(1, out->num_columns());
ASSERT_EQ(100, out->num_rows());
ASSERT_EQ(SMALL_SIZE, out->num_rows());
std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data();
ASSERT_EQ(1, chunked_array->num_chunks());
Int64Array* out_array = static_cast<Int64Array*>(chunked_array->chunk(0).get());
for (size_t i = 0; i < values.size(); i++) {
EXPECT_EQ(values[i], out_array->raw_data()[i]);
}
ExpectArray<typename TypeParam::c_type>(values.data(), chunked_array->chunk(0).get());
}
TEST_F(TestParquetIO, SingleColumnInt64Write) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<Int64Type>(100, 128);
TYPED_TEST(TestParquetIO, SingleColumnRequiredWrite) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<TypeParam>(SMALL_SIZE, 128);
std::shared_ptr<GroupNode> schema =
MakeSchema(ParquetType::INT64, Repetition::REQUIRED);
FileWriter writer(default_memory_pool(), MakeWriter(schema));
this->MakeSchema(test_traits<TypeParam>::parquet_enum, Repetition::REQUIRED);
FileWriter writer(default_memory_pool(), this->MakeWriter(schema));
ASSERT_NO_THROW(ASSERT_OK(writer.NewRowGroup(values->length())));
ASSERT_NO_THROW(ASSERT_OK(writer.WriteFlatColumnChunk(values.get())));
ASSERT_NO_THROW(ASSERT_OK(writer.Close()));
std::shared_ptr<Array> out;
ReadSingleColumnFile(ReaderFromSink(), &out);
this->ReadSingleColumnFile(this->ReaderFromSink(), &out);
ASSERT_TRUE(values->Equals(out));
}
TEST_F(TestParquetIO, SingleColumnTableInt64Write) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<Int64Type>(100, 128);
TYPED_TEST(TestParquetIO, SingleColumnTableRequiredWrite) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<TypeParam>(SMALL_SIZE, 128);
std::shared_ptr<Table> table = MakeSimpleTable(values, false);
sink_ = std::make_shared<InMemoryOutputStream>();
this->sink_ = std::make_shared<InMemoryOutputStream>();
ASSERT_NO_THROW(ASSERT_OK(
WriteFlatTable(table.get(), default_memory_pool(), sink_, values->length())));
WriteFlatTable(table.get(), default_memory_pool(), this->sink_, values->length())));
std::shared_ptr<Table> out;
ReadTableFromFile(ReaderFromSink(), &out);
this->ReadTableFromFile(this->ReaderFromSink(), &out);
ASSERT_EQ(1, out->num_columns());
ASSERT_EQ(100, out->num_rows());
@@ -238,47 +238,48 @@ TEST_F(TestParquetIO, SingleColumnTableInt64Write) {
ASSERT_TRUE(values->Equals(chunked_array->chunk(0)));
}
TEST_F(TestParquetIO, SingleColumnDoubleReadWrite) {
TYPED_TEST(TestParquetIO, SingleColumnOptionalReadWrite) {
// This also tests max_definition_level = 1
std::shared_ptr<PrimitiveArray> values = NullableArray<DoubleType>(100, 128, 10);
std::shared_ptr<PrimitiveArray> values = NullableArray<TypeParam>(SMALL_SIZE, 128, 10);
std::shared_ptr<GroupNode> schema =
MakeSchema(ParquetType::DOUBLE, Repetition::OPTIONAL);
FileWriter writer(default_memory_pool(), MakeWriter(schema));
this->MakeSchema(test_traits<TypeParam>::parquet_enum, Repetition::OPTIONAL);
FileWriter writer(default_memory_pool(), this->MakeWriter(schema));
ASSERT_NO_THROW(ASSERT_OK(writer.NewRowGroup(values->length())));
ASSERT_NO_THROW(ASSERT_OK(writer.WriteFlatColumnChunk(values.get())));
ASSERT_NO_THROW(ASSERT_OK(writer.Close()));
std::shared_ptr<Array> out;
ReadSingleColumnFile(ReaderFromSink(), &out);
this->ReadSingleColumnFile(this->ReaderFromSink(), &out);
ASSERT_TRUE(values->Equals(out));
}
TEST_F(TestParquetIO, SingleColumnTableDoubleReadWrite) {
TYPED_TEST(TestParquetIO, SingleColumnTableOptionalReadWrite) {
// This also tests max_definition_level = 1
std::shared_ptr<PrimitiveArray> values = NullableArray<DoubleType>(100, 128, 10);
std::shared_ptr<PrimitiveArray> values = NullableArray<TypeParam>(SMALL_SIZE, 128, 10);
std::shared_ptr<Table> table = MakeSimpleTable(values, true);
sink_ = std::make_shared<InMemoryOutputStream>();
this->sink_ = std::make_shared<InMemoryOutputStream>();
ASSERT_NO_THROW(ASSERT_OK(
WriteFlatTable(table.get(), default_memory_pool(), sink_, values->length())));
WriteFlatTable(table.get(), default_memory_pool(), this->sink_, values->length())));
std::shared_ptr<Table> out;
ReadTableFromFile(ReaderFromSink(), &out);
this->ReadTableFromFile(this->ReaderFromSink(), &out);
ASSERT_EQ(1, out->num_columns());
ASSERT_EQ(100, out->num_rows());
ASSERT_EQ(SMALL_SIZE, out->num_rows());
std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data();
ASSERT_EQ(1, chunked_array->num_chunks());
ASSERT_TRUE(values->Equals(chunked_array->chunk(0)));
}
TEST_F(TestParquetIO, SingleColumnInt64ChunkedWrite) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<Int64Type>(100, 128);
std::shared_ptr<PrimitiveArray> values_chunk = NonNullArray<Int64Type>(25, 128);
TYPED_TEST(TestParquetIO, SingleColumnIntRequiredChunkedWrite) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<TypeParam>(SMALL_SIZE, 128);
std::shared_ptr<PrimitiveArray> values_chunk =
NonNullArray<TypeParam>(SMALL_SIZE / 4, 128);
std::shared_ptr<GroupNode> schema =
MakeSchema(ParquetType::INT64, Repetition::REQUIRED);
FileWriter writer(default_memory_pool(), MakeWriter(schema));
this->MakeSchema(test_traits<TypeParam>::parquet_enum, Repetition::REQUIRED);
FileWriter writer(default_memory_pool(), this->MakeWriter(schema));
for (int i = 0; i < 4; i++) {
ASSERT_NO_THROW(ASSERT_OK(writer.NewRowGroup(values_chunk->length())));
ASSERT_NO_THROW(ASSERT_OK(writer.WriteFlatColumnChunk(values_chunk.get())));
@@ -286,36 +287,37 @@ TEST_F(TestParquetIO, SingleColumnInt64ChunkedWrite) {
ASSERT_NO_THROW(ASSERT_OK(writer.Close()));
std::shared_ptr<Array> out;
ReadSingleColumnFile(ReaderFromSink(), &out);
this->ReadSingleColumnFile(this->ReaderFromSink(), &out);
ASSERT_TRUE(values->Equals(out));
}
TEST_F(TestParquetIO, SingleColumnTableInt64ChunkedWrite) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<Int64Type>(1000, 128);
TYPED_TEST(TestParquetIO, SingleColumnTableRequiredChunkedWrite) {
std::shared_ptr<PrimitiveArray> values = NonNullArray<TypeParam>(LARGE_SIZE, 128);
std::shared_ptr<Table> table = MakeSimpleTable(values, false);
sink_ = std::make_shared<InMemoryOutputStream>();
this->sink_ = std::make_shared<InMemoryOutputStream>();
ASSERT_NO_THROW(
ASSERT_OK(WriteFlatTable(table.get(), default_memory_pool(), sink_, 512)));
ASSERT_OK(WriteFlatTable(table.get(), default_memory_pool(), this->sink_, 512)));
std::shared_ptr<Table> out;
ReadTableFromFile(ReaderFromSink(), &out);
this->ReadTableFromFile(this->ReaderFromSink(), &out);
ASSERT_EQ(1, out->num_columns());
ASSERT_EQ(1000, out->num_rows());
ASSERT_EQ(LARGE_SIZE, out->num_rows());
std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data();
ASSERT_EQ(1, chunked_array->num_chunks());
ASSERT_TRUE(values->Equals(chunked_array->chunk(0)));
}
TEST_F(TestParquetIO, SingleColumnDoubleChunkedWrite) {
std::shared_ptr<PrimitiveArray> values = NullableArray<DoubleType>(100, 128, 10);
TYPED_TEST(TestParquetIO, SingleColumnOptionalChunkedWrite) {
std::shared_ptr<PrimitiveArray> values = NullableArray<TypeParam>(SMALL_SIZE, 128, 10);
std::shared_ptr<PrimitiveArray> values_chunk_nulls =
NullableArray<DoubleType>(25, 128, 10);
std::shared_ptr<PrimitiveArray> values_chunk = NullableArray<DoubleType>(25, 128, 0);
NullableArray<TypeParam>(SMALL_SIZE / 4, 128, 10);
std::shared_ptr<PrimitiveArray> values_chunk =
NullableArray<TypeParam>(SMALL_SIZE / 4, 128, 0);
std::shared_ptr<GroupNode> schema =
MakeSchema(ParquetType::DOUBLE, Repetition::OPTIONAL);
FileWriter writer(default_memory_pool(), MakeWriter(schema));
this->MakeSchema(test_traits<TypeParam>::parquet_enum, Repetition::OPTIONAL);
FileWriter writer(default_memory_pool(), this->MakeWriter(schema));
ASSERT_NO_THROW(ASSERT_OK(writer.NewRowGroup(values_chunk_nulls->length())));
ASSERT_NO_THROW(ASSERT_OK(writer.WriteFlatColumnChunk(values_chunk_nulls.get())));
for (int i = 0; i < 3; i++) {
@@ -325,22 +327,22 @@ TEST_F(TestParquetIO, SingleColumnDoubleChunkedWrite) {
ASSERT_NO_THROW(ASSERT_OK(writer.Close()));
std::shared_ptr<Array> out;
ReadSingleColumnFile(ReaderFromSink(), &out);
this->ReadSingleColumnFile(this->ReaderFromSink(), &out);
ASSERT_TRUE(values->Equals(out));
}
TEST_F(TestParquetIO, SingleColumnTableDoubleChunkedWrite) {
TYPED_TEST(TestParquetIO, SingleColumnTableOptionalChunkedWrite) {
// This also tests max_definition_level = 1
std::shared_ptr<PrimitiveArray> values = NullableArray<DoubleType>(1000, 128, 100);
std::shared_ptr<PrimitiveArray> values = NullableArray<TypeParam>(LARGE_SIZE, 128, 100);
std::shared_ptr<Table> table = MakeSimpleTable(values, true);
sink_ = std::make_shared<InMemoryOutputStream>();
this->sink_ = std::make_shared<InMemoryOutputStream>();
ASSERT_NO_THROW(
ASSERT_OK(WriteFlatTable(table.get(), default_memory_pool(), sink_, 512)));
ASSERT_OK(WriteFlatTable(table.get(), default_memory_pool(), this->sink_, 512)));
std::shared_ptr<Table> out;
ReadTableFromFile(ReaderFromSink(), &out);
this->ReadTableFromFile(this->ReaderFromSink(), &out);
ASSERT_EQ(1, out->num_columns());
ASSERT_EQ(1000, out->num_rows());
ASSERT_EQ(LARGE_SIZE, out->num_rows());
std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data();
ASSERT_EQ(1, chunked_array->num_chunks());