| // 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. |
| |
| #ifdef _MSC_VER |
| #pragma warning(push) |
| // Disable forcing value to bool warnings |
| #pragma warning(disable : 4800) |
| #endif |
| |
| #include "gtest/gtest.h" |
| |
| #include <sstream> |
| |
| #include "parquet/api/reader.h" |
| #include "parquet/api/writer.h" |
| |
| #include "parquet/arrow/reader.h" |
| #include "parquet/arrow/schema.h" |
| #include "parquet/arrow/test-util.h" |
| #include "parquet/arrow/writer.h" |
| |
| #include "parquet/file/writer.h" |
| |
| #include "arrow/api.h" |
| #include "arrow/test-util.h" |
| |
| using arrow::Array; |
| using arrow::ArrayVisitor; |
| using arrow::Buffer; |
| using arrow::ChunkedArray; |
| using arrow::Column; |
| using arrow::EncodeArrayToDictionary; |
| using arrow::ListArray; |
| using arrow::PoolBuffer; |
| using arrow::PrimitiveArray; |
| using arrow::Status; |
| using arrow::Table; |
| using arrow::TimeUnit; |
| using arrow::default_memory_pool; |
| using arrow::io::BufferReader; |
| |
| using arrow::test::randint; |
| using arrow::test::random_is_valid; |
| |
| using ArrowId = ::arrow::Type; |
| using ParquetType = parquet::Type; |
| using parquet::schema::GroupNode; |
| using parquet::schema::NodePtr; |
| using parquet::schema::PrimitiveNode; |
| using parquet::arrow::FromParquetSchema; |
| |
| using ColumnVector = std::vector<std::shared_ptr<arrow::Column>>; |
| |
| namespace parquet { |
| namespace arrow { |
| |
| const int SMALL_SIZE = 100; |
| const int LARGE_SIZE = 10000; |
| |
| constexpr uint32_t kDefaultSeed = 0; |
| |
| LogicalType::type get_logical_type(const ::arrow::DataType& type) { |
| switch (type.id()) { |
| case ArrowId::UINT8: |
| return LogicalType::UINT_8; |
| case ArrowId::INT8: |
| return LogicalType::INT_8; |
| case ArrowId::UINT16: |
| return LogicalType::UINT_16; |
| case ArrowId::INT16: |
| return LogicalType::INT_16; |
| case ArrowId::UINT32: |
| return LogicalType::UINT_32; |
| case ArrowId::INT32: |
| return LogicalType::INT_32; |
| case ArrowId::UINT64: |
| return LogicalType::UINT_64; |
| case ArrowId::INT64: |
| return LogicalType::INT_64; |
| case ArrowId::STRING: |
| return LogicalType::UTF8; |
| case ArrowId::DATE32: |
| return LogicalType::DATE; |
| case ArrowId::DATE64: |
| return LogicalType::DATE; |
| case ArrowId::TIMESTAMP: { |
| const auto& ts_type = static_cast<const ::arrow::TimestampType&>(type); |
| switch (ts_type.unit()) { |
| case TimeUnit::MILLI: |
| return LogicalType::TIMESTAMP_MILLIS; |
| case TimeUnit::MICRO: |
| return LogicalType::TIMESTAMP_MICROS; |
| default: |
| DCHECK(false) << "Only MILLI and MICRO units supported for Arrow timestamps " |
| "with Parquet."; |
| } |
| } |
| case ArrowId::TIME32: |
| return LogicalType::TIME_MILLIS; |
| case ArrowId::TIME64: |
| return LogicalType::TIME_MICROS; |
| case ArrowId::DICTIONARY: { |
| const ::arrow::DictionaryType& dict_type = |
| static_cast<const ::arrow::DictionaryType&>(type); |
| return get_logical_type(*dict_type.dictionary()->type()); |
| } |
| default: |
| break; |
| } |
| return LogicalType::NONE; |
| } |
| |
| ParquetType::type get_physical_type(const ::arrow::DataType& type) { |
| switch (type.id()) { |
| case ArrowId::BOOL: |
| return ParquetType::BOOLEAN; |
| case ArrowId::UINT8: |
| case ArrowId::INT8: |
| case ArrowId::UINT16: |
| case ArrowId::INT16: |
| case ArrowId::UINT32: |
| case ArrowId::INT32: |
| return ParquetType::INT32; |
| case ArrowId::UINT64: |
| case ArrowId::INT64: |
| return ParquetType::INT64; |
| case ArrowId::FLOAT: |
| return ParquetType::FLOAT; |
| case ArrowId::DOUBLE: |
| return ParquetType::DOUBLE; |
| case ArrowId::BINARY: |
| return ParquetType::BYTE_ARRAY; |
| case ArrowId::STRING: |
| return ParquetType::BYTE_ARRAY; |
| case ArrowId::FIXED_SIZE_BINARY: |
| return ParquetType::FIXED_LEN_BYTE_ARRAY; |
| case ArrowId::DATE32: |
| return ParquetType::INT32; |
| case ArrowId::DATE64: |
| // Convert to date32 internally |
| return ParquetType::INT32; |
| case ArrowId::TIME32: |
| return ParquetType::INT32; |
| case ArrowId::TIME64: |
| return ParquetType::INT64; |
| case ArrowId::TIMESTAMP: |
| return ParquetType::INT64; |
| case ArrowId::DICTIONARY: { |
| const ::arrow::DictionaryType& dict_type = |
| static_cast<const ::arrow::DictionaryType&>(type); |
| return get_physical_type(*dict_type.dictionary()->type()); |
| } |
| default: |
| break; |
| } |
| DCHECK(false) << "cannot reach this code"; |
| return ParquetType::INT32; |
| } |
| |
| template <typename TestType> |
| struct test_traits {}; |
| |
| template <> |
| struct test_traits<::arrow::BooleanType> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::BOOLEAN; |
| static uint8_t const value; |
| }; |
| |
| const uint8_t test_traits<::arrow::BooleanType>::value(1); |
| |
| template <> |
| struct test_traits<::arrow::UInt8Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT32; |
| static uint8_t const value; |
| }; |
| |
| const uint8_t test_traits<::arrow::UInt8Type>::value(64); |
| |
| template <> |
| struct test_traits<::arrow::Int8Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT32; |
| static int8_t const value; |
| }; |
| |
| const int8_t test_traits<::arrow::Int8Type>::value(-64); |
| |
| template <> |
| struct test_traits<::arrow::UInt16Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT32; |
| static uint16_t const value; |
| }; |
| |
| const uint16_t test_traits<::arrow::UInt16Type>::value(1024); |
| |
| template <> |
| struct test_traits<::arrow::Int16Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT32; |
| static int16_t const value; |
| }; |
| |
| const int16_t test_traits<::arrow::Int16Type>::value(-1024); |
| |
| template <> |
| struct test_traits<::arrow::UInt32Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT32; |
| static uint32_t const value; |
| }; |
| |
| const uint32_t test_traits<::arrow::UInt32Type>::value(1024); |
| |
| template <> |
| struct test_traits<::arrow::Int32Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT32; |
| static int32_t const value; |
| }; |
| |
| const int32_t test_traits<::arrow::Int32Type>::value(-1024); |
| |
| template <> |
| struct test_traits<::arrow::UInt64Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT64; |
| static uint64_t const value; |
| }; |
| |
| const uint64_t test_traits<::arrow::UInt64Type>::value(1024); |
| |
| template <> |
| struct test_traits<::arrow::Int64Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT64; |
| static int64_t const value; |
| }; |
| |
| const int64_t test_traits<::arrow::Int64Type>::value(-1024); |
| |
| template <> |
| struct test_traits<::arrow::TimestampType> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT64; |
| static int64_t const value; |
| }; |
| |
| const int64_t test_traits<::arrow::TimestampType>::value(14695634030000); |
| |
| template <> |
| struct test_traits<::arrow::Date32Type> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::INT32; |
| static int32_t const value; |
| }; |
| |
| const int32_t test_traits<::arrow::Date32Type>::value(170000); |
| |
| template <> |
| struct test_traits<::arrow::FloatType> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::FLOAT; |
| static float const value; |
| }; |
| |
| const float test_traits<::arrow::FloatType>::value(2.1f); |
| |
| template <> |
| struct test_traits<::arrow::DoubleType> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::DOUBLE; |
| static double const value; |
| }; |
| |
| const double test_traits<::arrow::DoubleType>::value(4.2); |
| |
| template <> |
| struct test_traits<::arrow::StringType> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::BYTE_ARRAY; |
| static std::string const value; |
| }; |
| |
| template <> |
| struct test_traits<::arrow::BinaryType> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::BYTE_ARRAY; |
| static std::string const value; |
| }; |
| |
| template <> |
| struct test_traits<::arrow::FixedSizeBinaryType> { |
| static constexpr ParquetType::type parquet_enum = ParquetType::FIXED_LEN_BYTE_ARRAY; |
| static std::string const value; |
| }; |
| |
| const std::string test_traits<::arrow::StringType>::value("Test"); // NOLINT |
| const std::string test_traits<::arrow::BinaryType>::value("\x00\x01\x02\x03"); // NOLINT |
| const std::string test_traits<::arrow::FixedSizeBinaryType>::value("Fixed"); // NOLINT |
| template <typename T> |
| using ParquetDataType = DataType<test_traits<T>::parquet_enum>; |
| |
| template <typename T> |
| using ParquetWriter = TypedColumnWriter<ParquetDataType<T>>; |
| |
| void WriteTableToBuffer(const std::shared_ptr<Table>& table, int num_threads, |
| int64_t row_group_size, |
| const std::shared_ptr<ArrowWriterProperties>& arrow_properties, |
| std::shared_ptr<Buffer>* out) { |
| auto sink = std::make_shared<InMemoryOutputStream>(); |
| |
| ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), sink, |
| row_group_size, default_writer_properties(), |
| arrow_properties)); |
| *out = sink->GetBuffer(); |
| } |
| |
| void DoSimpleRoundtrip(const std::shared_ptr<Table>& table, int num_threads, |
| int64_t row_group_size, const std::vector<int>& column_subset, |
| std::shared_ptr<Table>* out, |
| const std::shared_ptr<ArrowWriterProperties>& arrow_properties = |
| default_arrow_writer_properties()) { |
| std::shared_ptr<Buffer> buffer; |
| WriteTableToBuffer(table, num_threads, row_group_size, arrow_properties, &buffer); |
| |
| std::unique_ptr<FileReader> reader; |
| ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer), |
| ::arrow::default_memory_pool(), |
| ::parquet::default_reader_properties(), nullptr, &reader)); |
| |
| reader->set_num_threads(num_threads); |
| |
| if (column_subset.size() > 0) { |
| ASSERT_OK_NO_THROW(reader->ReadTable(column_subset, out)); |
| } else { |
| // Read everything |
| ASSERT_OK_NO_THROW(reader->ReadTable(out)); |
| } |
| } |
| |
| static std::shared_ptr<GroupNode> MakeSimpleSchema(const ::arrow::DataType& type, |
| Repetition::type repetition) { |
| int byte_width; |
| // Decimal is not implemented yet. |
| switch (type.id()) { |
| case ::arrow::Type::DICTIONARY: { |
| const ::arrow::DictionaryType& dict_type = |
| static_cast<const ::arrow::DictionaryType&>(type); |
| const ::arrow::DataType& values_type = *dict_type.dictionary()->type(); |
| if (values_type.id() == ::arrow::Type::FIXED_SIZE_BINARY) { |
| byte_width = |
| static_cast<const ::arrow::FixedSizeBinaryType&>(values_type).byte_width(); |
| } else { |
| byte_width = -1; |
| } |
| } break; |
| case ::arrow::Type::FIXED_SIZE_BINARY: |
| byte_width = static_cast<const ::arrow::FixedSizeBinaryType&>(type).byte_width(); |
| break; |
| default: |
| byte_width = -1; |
| } |
| auto pnode = PrimitiveNode::Make("column1", repetition, get_physical_type(type), |
| get_logical_type(type), byte_width); |
| NodePtr node_ = |
| GroupNode::Make("schema", Repetition::REQUIRED, std::vector<NodePtr>({pnode})); |
| return std::static_pointer_cast<GroupNode>(node_); |
| } |
| |
| namespace internal { |
| |
| void AssertArraysEqual(const Array &expected, const Array &actual) { |
| if (!actual.Equals(expected)) { |
| std::stringstream pp_result; |
| std::stringstream pp_expected; |
| |
| EXPECT_OK(::arrow::PrettyPrint(actual, 0, &pp_result)); |
| EXPECT_OK(::arrow::PrettyPrint(expected, 0, &pp_expected)); |
| FAIL() << "Got: \n" << pp_result.str() << "\nExpected: \n" << pp_expected.str(); |
| } |
| } |
| |
| } // namespace internal |
| |
| void AssertChunkedEqual(const ChunkedArray& expected, const ChunkedArray& actual) { |
| ASSERT_EQ(expected.num_chunks(), actual.num_chunks()) << "# chunks unequal"; |
| if (!actual.Equals(expected)) { |
| std::stringstream pp_result; |
| std::stringstream pp_expected; |
| |
| for (int i = 0; i < actual.num_chunks(); ++i) { |
| auto c1 = actual.chunk(i); |
| auto c2 = expected.chunk(i); |
| if (!c1->Equals(*c2)) { |
| EXPECT_OK(::arrow::PrettyPrint(*c1, 0, &pp_result)); |
| EXPECT_OK(::arrow::PrettyPrint(*c2, 0, &pp_expected)); |
| FAIL() << "Chunk " << i << " Got: " << pp_result.str() |
| << "\nExpected: " << pp_expected.str(); |
| } |
| } |
| } |
| } |
| |
| void AssertTablesEqual(const Table& expected, const Table& actual) { |
| ASSERT_EQ(expected.num_columns(), actual.num_columns()); |
| |
| for (int i = 0; i < actual.num_columns(); ++i) { |
| AssertChunkedEqual(*expected.column(i)->data(), *actual.column(i)->data()); |
| } |
| ASSERT_TRUE(actual.Equals(expected)); |
| } |
| |
| template <typename TestType> |
| class TestParquetIO : public ::testing::Test { |
| public: |
| virtual void SetUp() {} |
| |
| std::unique_ptr<ParquetFileWriter> MakeWriter( |
| const std::shared_ptr<GroupNode>& schema) { |
| sink_ = std::make_shared<InMemoryOutputStream>(); |
| return ParquetFileWriter::Open(sink_, schema); |
| } |
| |
| void ReaderFromSink(std::unique_ptr<FileReader>* out) { |
| std::shared_ptr<Buffer> buffer = sink_->GetBuffer(); |
| ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer), |
| ::arrow::default_memory_pool(), |
| ::parquet::default_reader_properties(), nullptr, out)); |
| } |
| |
| void ReadSingleColumnFile(std::unique_ptr<FileReader> file_reader, |
| std::shared_ptr<Array>* out) { |
| std::unique_ptr<ColumnReader> column_reader; |
| ASSERT_OK_NO_THROW(file_reader->GetColumn(0, &column_reader)); |
| ASSERT_NE(nullptr, column_reader.get()); |
| |
| ASSERT_OK(column_reader->NextBatch(SMALL_SIZE, out)); |
| ASSERT_NE(nullptr, out->get()); |
| } |
| |
| void ReadAndCheckSingleColumnFile(const Array& values) { |
| std::shared_ptr<Array> out; |
| |
| std::unique_ptr<FileReader> reader; |
| ReaderFromSink(&reader); |
| ReadSingleColumnFile(std::move(reader), &out); |
| |
| internal::AssertArraysEqual(values, *out); |
| } |
| |
| void ReadTableFromFile(std::unique_ptr<FileReader> reader, |
| std::shared_ptr<Table>* out) { |
| ASSERT_OK_NO_THROW(reader->ReadTable(out)); |
| auto key_value_metadata = |
| reader->parquet_reader()->metadata()->key_value_metadata().get(); |
| ASSERT_EQ(nullptr, key_value_metadata); |
| ASSERT_NE(nullptr, out->get()); |
| } |
| |
| void PrepareListTable(int64_t size, bool nullable_lists, bool nullable_elements, |
| int64_t null_count, std::shared_ptr<Table>* out) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NullableArray<TestType>(size * size, nullable_elements ? null_count : 0, |
| kDefaultSeed, &values)); |
| // Also test that slice offsets are respected |
| values = values->Slice(5, values->length() - 5); |
| std::shared_ptr<ListArray> lists; |
| ASSERT_OK(MakeListArray(values, size, nullable_lists ? null_count : 0, |
| nullable_elements, &lists)); |
| *out = MakeSimpleTable(lists->Slice(3, size - 6), nullable_lists); |
| } |
| |
| void PrepareListOfListTable(int64_t size, bool nullable_parent_lists, |
| bool nullable_lists, bool nullable_elements, |
| int64_t null_count, std::shared_ptr<Table>* out) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NullableArray<TestType>(size * 6, nullable_elements ? null_count : 0, |
| kDefaultSeed, &values)); |
| std::shared_ptr<ListArray> lists; |
| ASSERT_OK(MakeListArray(values, size * 3, nullable_lists ? null_count : 0, |
| nullable_elements, &lists)); |
| std::shared_ptr<ListArray> parent_lists; |
| ASSERT_OK(MakeListArray(lists, size, nullable_parent_lists ? null_count : 0, |
| nullable_lists, &parent_lists)); |
| *out = MakeSimpleTable(parent_lists, nullable_parent_lists); |
| } |
| |
| void ReadAndCheckSingleColumnTable(const std::shared_ptr<Array>& values) { |
| std::shared_ptr<::arrow::Table> out; |
| std::unique_ptr<FileReader> reader; |
| ReaderFromSink(&reader); |
| ReadTableFromFile(std::move(reader), &out); |
| ASSERT_EQ(1, out->num_columns()); |
| ASSERT_EQ(values->length(), out->num_rows()); |
| |
| std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data(); |
| ASSERT_EQ(1, chunked_array->num_chunks()); |
| auto result = chunked_array->chunk(0); |
| |
| internal::AssertArraysEqual(*values, *result); |
| } |
| |
| void CheckRoundTrip(const std::shared_ptr<Table>& table) { |
| std::shared_ptr<Table> result; |
| DoSimpleRoundtrip(table, 1, table->num_rows(), {}, &result); |
| |
| AssertTablesEqual(*table, *result); |
| } |
| |
| template <typename ArrayType> |
| void WriteColumn(const std::shared_ptr<GroupNode>& schema, |
| const std::shared_ptr<ArrayType>& values) { |
| FileWriter writer(::arrow::default_memory_pool(), MakeWriter(schema)); |
| ASSERT_OK_NO_THROW(writer.NewRowGroup(values->length())); |
| ASSERT_OK_NO_THROW(writer.WriteColumnChunk(*values)); |
| ASSERT_OK_NO_THROW(writer.Close()); |
| } |
| |
| std::shared_ptr<InMemoryOutputStream> sink_; |
| }; |
| |
| // We have separate tests for UInt32Type as this is currently the only type |
| // where a roundtrip does not yield the identical Array structure. |
| // There we write an UInt32 Array but receive an Int64 Array as result for |
| // Parquet version 1.0. |
| |
| typedef ::testing::Types<::arrow::BooleanType, ::arrow::UInt8Type, ::arrow::Int8Type, |
| ::arrow::UInt16Type, ::arrow::Int16Type, ::arrow::Int32Type, |
| ::arrow::UInt64Type, ::arrow::Int64Type, ::arrow::Date32Type, |
| ::arrow::FloatType, ::arrow::DoubleType, ::arrow::StringType, |
| ::arrow::BinaryType, ::arrow::FixedSizeBinaryType> |
| TestTypes; |
| |
| TYPED_TEST_CASE(TestParquetIO, TestTypes); |
| |
| TYPED_TEST(TestParquetIO, SingleColumnRequiredWrite) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NonNullArray<TypeParam>(SMALL_SIZE, &values)); |
| |
| std::shared_ptr<GroupNode> schema = |
| MakeSimpleSchema(*values->type(), Repetition::REQUIRED); |
| this->WriteColumn(schema, values); |
| |
| this->ReadAndCheckSingleColumnFile(*values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnTableRequiredWrite) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NonNullArray<TypeParam>(SMALL_SIZE, &values)); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, false); |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), this->sink_, |
| values->length(), default_writer_properties())); |
| |
| std::shared_ptr<Table> out; |
| std::unique_ptr<FileReader> reader; |
| this->ReaderFromSink(&reader); |
| this->ReadTableFromFile(std::move(reader), &out); |
| ASSERT_EQ(1, out->num_columns()); |
| ASSERT_EQ(100, out->num_rows()); |
| |
| std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data(); |
| ASSERT_EQ(1, chunked_array->num_chunks()); |
| |
| internal::AssertArraysEqual(*values, *chunked_array->chunk(0)); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnOptionalReadWrite) { |
| // This also tests max_definition_level = 1 |
| std::shared_ptr<Array> values; |
| |
| ASSERT_OK(NullableArray<TypeParam>(SMALL_SIZE, 10, kDefaultSeed, &values)); |
| |
| std::shared_ptr<GroupNode> schema = |
| MakeSimpleSchema(*values->type(), Repetition::OPTIONAL); |
| this->WriteColumn(schema, values); |
| |
| this->ReadAndCheckSingleColumnFile(*values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnOptionalDictionaryWrite) { |
| // Skip tests for BOOL as we don't create dictionaries for it. |
| if (TypeParam::type_id == ::arrow::Type::BOOL) { |
| return; |
| } |
| |
| std::shared_ptr<Array> values; |
| |
| ASSERT_OK(NullableArray<TypeParam>(SMALL_SIZE, 10, kDefaultSeed, &values)); |
| |
| std::shared_ptr<Array> dict_values; |
| ASSERT_OK(EncodeArrayToDictionary(*values, default_memory_pool(), &dict_values)); |
| std::shared_ptr<GroupNode> schema = |
| MakeSimpleSchema(*dict_values->type(), Repetition::OPTIONAL); |
| this->WriteColumn(schema, dict_values); |
| |
| this->ReadAndCheckSingleColumnFile(*values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnRequiredSliceWrite) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NonNullArray<TypeParam>(2 * SMALL_SIZE, &values)); |
| std::shared_ptr<GroupNode> schema = |
| MakeSimpleSchema(*values->type(), Repetition::REQUIRED); |
| |
| std::shared_ptr<Array> sliced_values = values->Slice(SMALL_SIZE / 2, SMALL_SIZE); |
| this->WriteColumn(schema, sliced_values); |
| this->ReadAndCheckSingleColumnFile(*sliced_values); |
| |
| // Slice offset 1 higher |
| sliced_values = values->Slice(SMALL_SIZE / 2 + 1, SMALL_SIZE); |
| this->WriteColumn(schema, sliced_values); |
| this->ReadAndCheckSingleColumnFile(*sliced_values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnOptionalSliceWrite) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NullableArray<TypeParam>(2 * SMALL_SIZE, SMALL_SIZE, kDefaultSeed, &values)); |
| std::shared_ptr<GroupNode> schema = |
| MakeSimpleSchema(*values->type(), Repetition::OPTIONAL); |
| |
| std::shared_ptr<Array> sliced_values = values->Slice(SMALL_SIZE / 2, SMALL_SIZE); |
| this->WriteColumn(schema, sliced_values); |
| this->ReadAndCheckSingleColumnFile(*sliced_values); |
| |
| // Slice offset 1 higher, thus different null bitmap. |
| sliced_values = values->Slice(SMALL_SIZE / 2 + 1, SMALL_SIZE); |
| this->WriteColumn(schema, sliced_values); |
| this->ReadAndCheckSingleColumnFile(*sliced_values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnTableOptionalReadWrite) { |
| // This also tests max_definition_level = 1 |
| std::shared_ptr<Array> values; |
| |
| ASSERT_OK(NullableArray<TypeParam>(SMALL_SIZE, 10, kDefaultSeed, &values)); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, true); |
| this->CheckRoundTrip(table); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleNullableListNullableColumnReadWrite) { |
| std::shared_ptr<Table> table; |
| this->PrepareListTable(SMALL_SIZE, true, true, 10, &table); |
| this->CheckRoundTrip(table); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleRequiredListNullableColumnReadWrite) { |
| std::shared_ptr<Table> table; |
| this->PrepareListTable(SMALL_SIZE, false, true, 10, &table); |
| this->CheckRoundTrip(table); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleNullableListRequiredColumnReadWrite) { |
| std::shared_ptr<Table> table; |
| this->PrepareListTable(SMALL_SIZE, true, false, 10, &table); |
| this->CheckRoundTrip(table); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleRequiredListRequiredColumnReadWrite) { |
| std::shared_ptr<Table> table; |
| this->PrepareListTable(SMALL_SIZE, false, false, 0, &table); |
| this->CheckRoundTrip(table); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleNullableListRequiredListRequiredColumnReadWrite) { |
| std::shared_ptr<Table> table; |
| this->PrepareListOfListTable(SMALL_SIZE, true, false, false, 0, &table); |
| this->CheckRoundTrip(table); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnRequiredChunkedWrite) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NonNullArray<TypeParam>(SMALL_SIZE, &values)); |
| int64_t chunk_size = values->length() / 4; |
| |
| std::shared_ptr<GroupNode> schema = |
| MakeSimpleSchema(*values->type(), Repetition::REQUIRED); |
| FileWriter writer(default_memory_pool(), this->MakeWriter(schema)); |
| for (int i = 0; i < 4; i++) { |
| ASSERT_OK_NO_THROW(writer.NewRowGroup(chunk_size)); |
| std::shared_ptr<Array> sliced_array = values->Slice(i * chunk_size, chunk_size); |
| ASSERT_OK_NO_THROW(writer.WriteColumnChunk(*sliced_array)); |
| } |
| ASSERT_OK_NO_THROW(writer.Close()); |
| |
| this->ReadAndCheckSingleColumnFile(*values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnTableRequiredChunkedWrite) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NonNullArray<TypeParam>(LARGE_SIZE, &values)); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, false); |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| ASSERT_OK_NO_THROW(WriteTable(*table, default_memory_pool(), this->sink_, 512, |
| default_writer_properties())); |
| |
| this->ReadAndCheckSingleColumnTable(values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnTableRequiredChunkedWriteArrowIO) { |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NonNullArray<TypeParam>(LARGE_SIZE, &values)); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, false); |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| auto buffer = std::make_shared<::arrow::PoolBuffer>(); |
| |
| { |
| // BufferOutputStream closed on gc |
| auto arrow_sink_ = std::make_shared<::arrow::io::BufferOutputStream>(buffer); |
| ASSERT_OK_NO_THROW(WriteTable(*table, default_memory_pool(), arrow_sink_, 512, |
| default_writer_properties())); |
| |
| // XXX: Remove this after ARROW-455 completed |
| ASSERT_OK(arrow_sink_->Close()); |
| } |
| |
| auto pbuffer = std::make_shared<Buffer>(buffer->data(), buffer->size()); |
| |
| auto source = std::make_shared<BufferReader>(pbuffer); |
| std::shared_ptr<::arrow::Table> out; |
| std::unique_ptr<FileReader> reader; |
| ASSERT_OK_NO_THROW(OpenFile(source, ::arrow::default_memory_pool(), &reader)); |
| this->ReadTableFromFile(std::move(reader), &out); |
| ASSERT_EQ(1, out->num_columns()); |
| ASSERT_EQ(values->length(), out->num_rows()); |
| |
| std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data(); |
| ASSERT_EQ(1, chunked_array->num_chunks()); |
| |
| internal::AssertArraysEqual(*values, *chunked_array->chunk(0)); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnOptionalChunkedWrite) { |
| int64_t chunk_size = SMALL_SIZE / 4; |
| std::shared_ptr<Array> values; |
| |
| ASSERT_OK(NullableArray<TypeParam>(SMALL_SIZE, 10, kDefaultSeed, &values)); |
| |
| std::shared_ptr<GroupNode> schema = |
| MakeSimpleSchema(*values->type(), Repetition::OPTIONAL); |
| FileWriter writer(::arrow::default_memory_pool(), this->MakeWriter(schema)); |
| for (int i = 0; i < 4; i++) { |
| ASSERT_OK_NO_THROW(writer.NewRowGroup(chunk_size)); |
| std::shared_ptr<Array> sliced_array = values->Slice(i * chunk_size, chunk_size); |
| ASSERT_OK_NO_THROW(writer.WriteColumnChunk(*sliced_array)); |
| } |
| ASSERT_OK_NO_THROW(writer.Close()); |
| |
| this->ReadAndCheckSingleColumnFile(*values); |
| } |
| |
| TYPED_TEST(TestParquetIO, SingleColumnTableOptionalChunkedWrite) { |
| // This also tests max_definition_level = 1 |
| std::shared_ptr<Array> values; |
| |
| ASSERT_OK(NullableArray<TypeParam>(LARGE_SIZE, 100, kDefaultSeed, &values)); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, true); |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), this->sink_, 512, |
| default_writer_properties())); |
| |
| this->ReadAndCheckSingleColumnTable(values); |
| } |
| |
| using TestInt96ParquetIO = TestParquetIO<::arrow::TimestampType>; |
| |
| TEST_F(TestInt96ParquetIO, ReadIntoTimestamp) { |
| // This test explicitly tests the conversion from an Impala-style timestamp |
| // to a nanoseconds-since-epoch one. |
| |
| // 2nd January 1970, 11:35min 145738543ns |
| Int96 day; |
| day.value[2] = UINT32_C(2440589); |
| int64_t seconds = (11 * 60 + 35) * 60; |
| *(reinterpret_cast<int64_t*>(&(day.value))) = |
| seconds * INT64_C(1000) * INT64_C(1000) * INT64_C(1000) + 145738543; |
| // Compute the corresponding nanosecond timestamp |
| struct tm datetime = {0}; |
| datetime.tm_year = 70; |
| datetime.tm_mon = 0; |
| datetime.tm_mday = 2; |
| datetime.tm_hour = 11; |
| datetime.tm_min = 35; |
| struct tm epoch = {0}; |
| epoch.tm_year = 70; |
| epoch.tm_mday = 1; |
| // Nanoseconds since the epoch |
| int64_t val = lrint(difftime(mktime(&datetime), mktime(&epoch))) * INT64_C(1000000000); |
| val += 145738543; |
| |
| std::vector<std::shared_ptr<schema::Node>> fields( |
| {schema::PrimitiveNode::Make("int96", Repetition::REQUIRED, ParquetType::INT96)}); |
| std::shared_ptr<schema::GroupNode> schema = std::static_pointer_cast<GroupNode>( |
| schema::GroupNode::Make("schema", Repetition::REQUIRED, fields)); |
| |
| // We cannot write this column with Arrow, so we have to use the plain parquet-cpp API |
| // to write an Int96 file. |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| auto writer = ParquetFileWriter::Open(this->sink_, schema); |
| RowGroupWriter* rg_writer = writer->AppendRowGroup(); |
| ColumnWriter* c_writer = rg_writer->NextColumn(); |
| auto typed_writer = dynamic_cast<TypedColumnWriter<Int96Type>*>(c_writer); |
| ASSERT_NE(typed_writer, nullptr); |
| typed_writer->WriteBatch(1, nullptr, nullptr, &day); |
| c_writer->Close(); |
| rg_writer->Close(); |
| writer->Close(); |
| |
| ::arrow::TimestampBuilder builder(::arrow::timestamp(TimeUnit::NANO), |
| ::arrow::default_memory_pool()); |
| ASSERT_OK(builder.Append(val)); |
| std::shared_ptr<Array> values; |
| ASSERT_OK(builder.Finish(&values)); |
| this->ReadAndCheckSingleColumnFile(*values); |
| } |
| |
| using TestUInt32ParquetIO = TestParquetIO<::arrow::UInt32Type>; |
| |
| TEST_F(TestUInt32ParquetIO, Parquet_2_0_Compability) { |
| // This also tests max_definition_level = 1 |
| std::shared_ptr<Array> values; |
| |
| ASSERT_OK(NullableArray<::arrow::UInt32Type>(LARGE_SIZE, 100, kDefaultSeed, &values)); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, true); |
| |
| // Parquet 2.0 roundtrip should yield an uint32_t column again |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| std::shared_ptr<::parquet::WriterProperties> properties = |
| ::parquet::WriterProperties::Builder() |
| .version(ParquetVersion::PARQUET_2_0) |
| ->build(); |
| ASSERT_OK_NO_THROW( |
| WriteTable(*table, default_memory_pool(), this->sink_, 512, properties)); |
| this->ReadAndCheckSingleColumnTable(values); |
| } |
| |
| TEST_F(TestUInt32ParquetIO, Parquet_1_0_Compability) { |
| // This also tests max_definition_level = 1 |
| std::shared_ptr<Array> arr; |
| ASSERT_OK(NullableArray<::arrow::UInt32Type>(LARGE_SIZE, 100, kDefaultSeed, &arr)); |
| |
| std::shared_ptr<::arrow::UInt32Array> values = |
| std::dynamic_pointer_cast<::arrow::UInt32Array>(arr); |
| |
| std::shared_ptr<Table> table = MakeSimpleTable(values, true); |
| |
| // Parquet 1.0 returns an int64_t column as there is no way to tell a Parquet 1.0 |
| // reader that a column is unsigned. |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| std::shared_ptr<::parquet::WriterProperties> properties = |
| ::parquet::WriterProperties::Builder() |
| .version(ParquetVersion::PARQUET_1_0) |
| ->build(); |
| ASSERT_OK_NO_THROW( |
| WriteTable(*table, ::arrow::default_memory_pool(), this->sink_, 512, properties)); |
| |
| std::shared_ptr<Array> expected_values; |
| std::shared_ptr<PoolBuffer> int64_data = |
| std::make_shared<PoolBuffer>(::arrow::default_memory_pool()); |
| { |
| ASSERT_OK(int64_data->Resize(sizeof(int64_t) * values->length())); |
| int64_t* int64_data_ptr = reinterpret_cast<int64_t*>(int64_data->mutable_data()); |
| const uint32_t* uint32_data_ptr = |
| reinterpret_cast<const uint32_t*>(values->values()->data()); |
| // std::copy might be faster but this is explicit on the casts) |
| for (int64_t i = 0; i < values->length(); i++) { |
| int64_data_ptr[i] = static_cast<int64_t>(uint32_data_ptr[i]); |
| } |
| } |
| |
| std::vector<std::shared_ptr<Buffer>> buffers{values->null_bitmap(), int64_data}; |
| auto arr_data = std::make_shared<::arrow::ArrayData>(::arrow::int64(), values->length(), |
| buffers, values->null_count()); |
| ASSERT_OK(MakeArray(arr_data, &expected_values)); |
| this->ReadAndCheckSingleColumnTable(expected_values); |
| } |
| |
| using TestStringParquetIO = TestParquetIO<::arrow::StringType>; |
| |
| TEST_F(TestStringParquetIO, EmptyStringColumnRequiredWrite) { |
| std::shared_ptr<Array> values; |
| ::arrow::StringBuilder builder; |
| for (size_t i = 0; i < SMALL_SIZE; i++) { |
| ASSERT_OK(builder.Append("")); |
| } |
| ASSERT_OK(builder.Finish(&values)); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, false); |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), this->sink_, |
| values->length(), default_writer_properties())); |
| |
| std::shared_ptr<Table> out; |
| std::unique_ptr<FileReader> reader; |
| this->ReaderFromSink(&reader); |
| this->ReadTableFromFile(std::move(reader), &out); |
| ASSERT_EQ(1, out->num_columns()); |
| ASSERT_EQ(100, out->num_rows()); |
| |
| std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data(); |
| ASSERT_EQ(1, chunked_array->num_chunks()); |
| |
| internal::AssertArraysEqual(*values, *chunked_array->chunk(0)); |
| } |
| |
| using TestNullParquetIO = TestParquetIO<::arrow::NullType>; |
| |
| TEST_F(TestNullParquetIO, NullColumn) { |
| std::shared_ptr<Array> values = std::make_shared<::arrow::NullArray>(SMALL_SIZE); |
| std::shared_ptr<Table> table = MakeSimpleTable(values, true); |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), this->sink_, |
| values->length(), default_writer_properties())); |
| |
| std::shared_ptr<Table> out; |
| std::unique_ptr<FileReader> reader; |
| this->ReaderFromSink(&reader); |
| this->ReadTableFromFile(std::move(reader), &out); |
| ASSERT_EQ(1, out->num_columns()); |
| ASSERT_EQ(100, out->num_rows()); |
| |
| std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data(); |
| ASSERT_EQ(1, chunked_array->num_chunks()); |
| |
| internal::AssertArraysEqual(*values, *chunked_array->chunk(0)); |
| } |
| |
| TEST_F(TestNullParquetIO, NullDictionaryColumn) { |
| std::shared_ptr<Array> values = std::make_shared<::arrow::NullArray>(0); |
| std::shared_ptr<Array> indices = |
| std::make_shared<::arrow::Int8Array>(SMALL_SIZE, nullptr, nullptr, SMALL_SIZE); |
| std::shared_ptr<::arrow::DictionaryType> dict_type = |
| std::make_shared<::arrow::DictionaryType>(::arrow::int8(), values); |
| std::shared_ptr<Array> dict_values = |
| std::make_shared<::arrow::DictionaryArray>(dict_type, indices); |
| std::shared_ptr<Table> table = MakeSimpleTable(dict_values, true); |
| this->sink_ = std::make_shared<InMemoryOutputStream>(); |
| ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), this->sink_, |
| dict_values->length(), default_writer_properties())); |
| |
| std::shared_ptr<Table> out; |
| std::unique_ptr<FileReader> reader; |
| this->ReaderFromSink(&reader); |
| this->ReadTableFromFile(std::move(reader), &out); |
| ASSERT_EQ(1, out->num_columns()); |
| ASSERT_EQ(100, out->num_rows()); |
| |
| std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data(); |
| ASSERT_EQ(1, chunked_array->num_chunks()); |
| |
| std::shared_ptr<Array> expected_values = |
| std::make_shared<::arrow::NullArray>(SMALL_SIZE); |
| internal::AssertArraysEqual(*expected_values, *chunked_array->chunk(0)); |
| } |
| |
| template <typename T> |
| using ParquetCDataType = typename ParquetDataType<T>::c_type; |
| |
| template <typename T> |
| struct c_type_trait { |
| using ArrowCType = typename T::c_type; |
| }; |
| |
| template <> |
| struct c_type_trait<::arrow::BooleanType> { |
| using ArrowCType = uint8_t; |
| }; |
| |
| template <typename TestType> |
| class TestPrimitiveParquetIO : public TestParquetIO<TestType> { |
| public: |
| typedef typename c_type_trait<TestType>::ArrowCType T; |
| |
| void MakeTestFile(std::vector<T>& values, int num_chunks, |
| std::unique_ptr<FileReader>* reader) { |
| TestType dummy; |
| |
| std::shared_ptr<GroupNode> schema = MakeSimpleSchema(dummy, Repetition::REQUIRED); |
| std::unique_ptr<ParquetFileWriter> file_writer = this->MakeWriter(schema); |
| size_t chunk_size = values.size() / num_chunks; |
| // Convert to Parquet's expected physical type |
| std::vector<uint8_t> values_buffer(sizeof(ParquetCDataType<TestType>) * |
| values.size()); |
| auto values_parquet = |
| reinterpret_cast<ParquetCDataType<TestType>*>(values_buffer.data()); |
| std::copy(values.cbegin(), values.cend(), values_parquet); |
| for (int i = 0; i < num_chunks; i++) { |
| auto row_group_writer = file_writer->AppendRowGroup(); |
| auto column_writer = |
| static_cast<ParquetWriter<TestType>*>(row_group_writer->NextColumn()); |
| ParquetCDataType<TestType>* data = values_parquet + i * chunk_size; |
| column_writer->WriteBatch(chunk_size, nullptr, nullptr, data); |
| column_writer->Close(); |
| row_group_writer->Close(); |
| } |
| file_writer->Close(); |
| this->ReaderFromSink(reader); |
| } |
| |
| void CheckSingleColumnRequiredTableRead(int num_chunks) { |
| std::vector<T> values(SMALL_SIZE, test_traits<TestType>::value); |
| std::unique_ptr<FileReader> file_reader; |
| ASSERT_NO_THROW(MakeTestFile(values, num_chunks, &file_reader)); |
| |
| std::shared_ptr<Table> out; |
| this->ReadTableFromFile(std::move(file_reader), &out); |
| ASSERT_EQ(1, out->num_columns()); |
| ASSERT_EQ(SMALL_SIZE, out->num_rows()); |
| |
| std::shared_ptr<ChunkedArray> chunked_array = out->column(0)->data(); |
| ASSERT_EQ(1, chunked_array->num_chunks()); |
| ExpectArrayT<TestType>(values.data(), chunked_array->chunk(0).get()); |
| } |
| |
| void CheckSingleColumnRequiredRead(int num_chunks) { |
| std::vector<T> values(SMALL_SIZE, test_traits<TestType>::value); |
| std::unique_ptr<FileReader> file_reader; |
| ASSERT_NO_THROW(MakeTestFile(values, num_chunks, &file_reader)); |
| |
| std::shared_ptr<Array> out; |
| this->ReadSingleColumnFile(std::move(file_reader), &out); |
| |
| ExpectArrayT<TestType>(values.data(), out.get()); |
| } |
| }; |
| |
| typedef ::testing::Types<::arrow::BooleanType, ::arrow::UInt8Type, ::arrow::Int8Type, |
| ::arrow::UInt16Type, ::arrow::Int16Type, ::arrow::UInt32Type, |
| ::arrow::Int32Type, ::arrow::UInt64Type, ::arrow::Int64Type, |
| ::arrow::FloatType, ::arrow::DoubleType> |
| PrimitiveTestTypes; |
| |
| TYPED_TEST_CASE(TestPrimitiveParquetIO, PrimitiveTestTypes); |
| |
| TYPED_TEST(TestPrimitiveParquetIO, SingleColumnRequiredRead) { |
| this->CheckSingleColumnRequiredRead(1); |
| } |
| |
| TYPED_TEST(TestPrimitiveParquetIO, SingleColumnRequiredTableRead) { |
| this->CheckSingleColumnRequiredTableRead(1); |
| } |
| |
| TYPED_TEST(TestPrimitiveParquetIO, SingleColumnRequiredChunkedRead) { |
| this->CheckSingleColumnRequiredRead(4); |
| } |
| |
| TYPED_TEST(TestPrimitiveParquetIO, SingleColumnRequiredChunkedTableRead) { |
| this->CheckSingleColumnRequiredTableRead(4); |
| } |
| |
| void MakeDateTimeTypesTable(std::shared_ptr<Table>* out, bool nanos_as_micros = false) { |
| using ::arrow::ArrayFromVector; |
| |
| std::vector<bool> is_valid = {true, true, true, false, true, true}; |
| |
| // These are only types that roundtrip without modification |
| auto f0 = field("f0", ::arrow::date32()); |
| auto f1 = field("f1", ::arrow::timestamp(TimeUnit::MILLI)); |
| auto f2 = field("f2", ::arrow::timestamp(TimeUnit::MICRO)); |
| std::shared_ptr<::arrow::Field> f3; |
| if (nanos_as_micros) { |
| f3 = field("f3", ::arrow::timestamp(TimeUnit::MICRO)); |
| } else { |
| f3 = field("f3", ::arrow::timestamp(TimeUnit::NANO)); |
| } |
| auto f4 = field("f4", ::arrow::time32(TimeUnit::MILLI)); |
| auto f5 = field("f5", ::arrow::time64(TimeUnit::MICRO)); |
| std::shared_ptr<::arrow::Schema> schema(new ::arrow::Schema({f0, f1, f2, f3, f4, f5})); |
| |
| std::vector<int32_t> t32_values = {1489269000, 1489270000, 1489271000, |
| 1489272000, 1489272000, 1489273000}; |
| std::vector<int64_t> t64_values = {1489269000000, 1489270000000, 1489271000000, |
| 1489272000000, 1489272000000, 1489273000000}; |
| std::vector<int64_t> t64_us_values = {1489269000, 1489270000, 1489271000, |
| 1489272000, 1489272000, 1489273000}; |
| |
| std::shared_ptr<Array> a0, a1, a2, a3, a4, a5; |
| ArrayFromVector<::arrow::Date32Type, int32_t>(f0->type(), is_valid, t32_values, &a0); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(f1->type(), is_valid, t64_values, &a1); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(f2->type(), is_valid, t64_values, &a2); |
| if (nanos_as_micros) { |
| ArrayFromVector<::arrow::TimestampType, int64_t>(f3->type(), is_valid, t64_us_values, |
| &a3); |
| } else { |
| ArrayFromVector<::arrow::TimestampType, int64_t>(f3->type(), is_valid, t64_values, |
| &a3); |
| } |
| ArrayFromVector<::arrow::Time32Type, int32_t>(f4->type(), is_valid, t32_values, &a4); |
| ArrayFromVector<::arrow::Time64Type, int64_t>(f5->type(), is_valid, t64_values, &a5); |
| |
| std::vector<std::shared_ptr<::arrow::Column>> columns = { |
| std::make_shared<Column>("f0", a0), std::make_shared<Column>("f1", a1), |
| std::make_shared<Column>("f2", a2), std::make_shared<Column>("f3", a3), |
| std::make_shared<Column>("f4", a4), std::make_shared<Column>("f5", a5)}; |
| *out = std::make_shared<::arrow::Table>(schema, columns); |
| } |
| |
| TEST(TestArrowReadWrite, DateTimeTypes) { |
| std::shared_ptr<Table> table; |
| MakeDateTimeTypesTable(&table); |
| |
| // Use deprecated INT96 type |
| std::shared_ptr<Table> result; |
| DoSimpleRoundtrip( |
| table, 1, table->num_rows(), {}, &result, |
| ArrowWriterProperties::Builder().enable_deprecated_int96_timestamps()->build()); |
| |
| AssertTablesEqual(*table, *result); |
| |
| // Cast nanaoseconds to microseconds and use INT64 physical type |
| DoSimpleRoundtrip(table, 1, table->num_rows(), {}, &result); |
| std::shared_ptr<Table> expected; |
| MakeDateTimeTypesTable(&table, true); |
| |
| AssertTablesEqual(*table, *result); |
| } |
| |
| TEST(TestArrowReadWrite, CoerceTimestamps) { |
| using ::arrow::ArrayFromVector; |
| using ::arrow::field; |
| |
| // PARQUET-1078, coerce Arrow timestamps to either TIMESTAMP_MILLIS or TIMESTAMP_MICROS |
| std::vector<bool> is_valid = {true, true, true, false, true, true}; |
| |
| auto t_s = ::arrow::timestamp(TimeUnit::SECOND); |
| auto t_ms = ::arrow::timestamp(TimeUnit::MILLI); |
| auto t_us = ::arrow::timestamp(TimeUnit::MICRO); |
| auto t_ns = ::arrow::timestamp(TimeUnit::NANO); |
| |
| std::vector<int64_t> s_values = {1489269, 1489270, 1489271, 1489272, 1489272, 1489273}; |
| std::vector<int64_t> ms_values = {1489269000, 1489270000, 1489271000, |
| 1489272001, 1489272000, 1489273000}; |
| std::vector<int64_t> us_values = {1489269000000, 1489270000000, 1489271000000, |
| 1489272000001, 1489272000000, 1489273000000}; |
| std::vector<int64_t> ns_values = {1489269000000000LL, 1489270000000000LL, |
| 1489271000000000LL, 1489272000000001LL, |
| 1489272000000000LL, 1489273000000000LL}; |
| |
| std::shared_ptr<Array> a_s, a_ms, a_us, a_ns; |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_s, is_valid, s_values, &a_s); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_ms, is_valid, ms_values, &a_ms); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_us, is_valid, us_values, &a_us); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_ns, is_valid, ns_values, &a_ns); |
| |
| // Input table, all data as is |
| auto s1 = std::shared_ptr<::arrow::Schema>( |
| new ::arrow::Schema({field("f_s", t_s), field("f_ms", t_ms), field("f_us", t_us), |
| field("f_ns", t_ns)})); |
| auto input = std::make_shared<::arrow::Table>( |
| s1, ColumnVector({std::make_shared<Column>("f_s", a_s), |
| std::make_shared<Column>("f_ms", a_ms), |
| std::make_shared<Column>("f_us", a_us), |
| std::make_shared<Column>("f_ns", a_ns)})); |
| |
| // Result when coercing to milliseconds |
| auto s2 = std::shared_ptr<::arrow::Schema>( |
| new ::arrow::Schema({field("f_s", t_ms), field("f_ms", t_ms), field("f_us", t_ms), |
| field("f_ns", t_ms)})); |
| auto ex_milli_result = std::make_shared<::arrow::Table>( |
| s2, ColumnVector({std::make_shared<Column>("f_s", a_ms), |
| std::make_shared<Column>("f_ms", a_ms), |
| std::make_shared<Column>("f_us", a_ms), |
| std::make_shared<Column>("f_ns", a_ms)})); |
| |
| // Result when coercing to microseconds |
| auto s3 = std::shared_ptr<::arrow::Schema>( |
| new ::arrow::Schema({field("f_s", t_us), field("f_ms", t_us), field("f_us", t_us), |
| field("f_ns", t_us)})); |
| auto ex_micro_result = std::make_shared<::arrow::Table>( |
| s3, ColumnVector({std::make_shared<Column>("f_s", a_us), |
| std::make_shared<Column>("f_ms", a_us), |
| std::make_shared<Column>("f_us", a_us), |
| std::make_shared<Column>("f_ns", a_us)})); |
| |
| std::shared_ptr<Table> milli_result; |
| DoSimpleRoundtrip( |
| input, 1, input->num_rows(), {}, &milli_result, |
| ArrowWriterProperties::Builder().coerce_timestamps(TimeUnit::MILLI)->build()); |
| |
| AssertTablesEqual(*ex_milli_result, *milli_result); |
| |
| std::shared_ptr<Table> micro_result; |
| DoSimpleRoundtrip( |
| input, 1, input->num_rows(), {}, µ_result, |
| ArrowWriterProperties::Builder().coerce_timestamps(TimeUnit::MICRO)->build()); |
| AssertTablesEqual(*ex_micro_result, *micro_result); |
| } |
| |
| TEST(TestArrowReadWrite, CoerceTimestampsLosePrecision) { |
| using ::arrow::ArrayFromVector; |
| using ::arrow::field; |
| |
| // PARQUET-1078, coerce Arrow timestamps to either TIMESTAMP_MILLIS or TIMESTAMP_MICROS |
| std::vector<bool> is_valid = {true, true, true, false, true, true}; |
| |
| auto t_s = ::arrow::timestamp(TimeUnit::SECOND); |
| auto t_ms = ::arrow::timestamp(TimeUnit::MILLI); |
| auto t_us = ::arrow::timestamp(TimeUnit::MICRO); |
| auto t_ns = ::arrow::timestamp(TimeUnit::NANO); |
| |
| std::vector<int64_t> s_values = {1489269, 1489270, 1489271, 1489272, 1489272, 1489273}; |
| std::vector<int64_t> ms_values = {1489269001, 1489270001, 1489271001, |
| 1489272001, 1489272001, 1489273001}; |
| std::vector<int64_t> us_values = {1489269000001, 1489270000001, 1489271000001, |
| 1489272000001, 1489272000001, 1489273000001}; |
| std::vector<int64_t> ns_values = {1489269000000001LL, 1489270000000001LL, |
| 1489271000000001LL, 1489272000000001LL, |
| 1489272000000001LL, 1489273000000001LL}; |
| |
| std::shared_ptr<Array> a_s, a_ms, a_us, a_ns; |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_s, is_valid, s_values, &a_s); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_ms, is_valid, ms_values, &a_ms); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_us, is_valid, us_values, &a_us); |
| ArrayFromVector<::arrow::TimestampType, int64_t>(t_ns, is_valid, ns_values, &a_ns); |
| |
| auto s1 = std::shared_ptr<::arrow::Schema>(new ::arrow::Schema({field("f_s", t_s)})); |
| auto s2 = std::shared_ptr<::arrow::Schema>(new ::arrow::Schema({field("f_ms", t_ms)})); |
| auto s3 = std::shared_ptr<::arrow::Schema>(new ::arrow::Schema({field("f_us", t_us)})); |
| auto s4 = std::shared_ptr<::arrow::Schema>(new ::arrow::Schema({field("f_ns", t_ns)})); |
| |
| auto c1 = std::make_shared<Column>("f_s", a_s); |
| auto c2 = std::make_shared<Column>("f_ms", a_ms); |
| auto c3 = std::make_shared<Column>("f_us", a_us); |
| auto c4 = std::make_shared<Column>("f_ns", a_ns); |
| |
| auto t1 = std::make_shared<::arrow::Table>(s1, ColumnVector({c1})); |
| auto t2 = std::make_shared<::arrow::Table>(s2, ColumnVector({c2})); |
| auto t3 = std::make_shared<::arrow::Table>(s3, ColumnVector({c3})); |
| auto t4 = std::make_shared<::arrow::Table>(s4, ColumnVector({c4})); |
| |
| auto sink = std::make_shared<InMemoryOutputStream>(); |
| |
| // OK to write to millis |
| auto coerce_millis = |
| (ArrowWriterProperties::Builder().coerce_timestamps(TimeUnit::MILLI)->build()); |
| ASSERT_OK_NO_THROW(WriteTable(*t1, ::arrow::default_memory_pool(), sink, 10, |
| default_writer_properties(), coerce_millis)); |
| ASSERT_OK_NO_THROW(WriteTable(*t2, ::arrow::default_memory_pool(), sink, 10, |
| default_writer_properties(), coerce_millis)); |
| |
| // Loss of precision |
| ASSERT_RAISES(Invalid, WriteTable(*t3, ::arrow::default_memory_pool(), sink, 10, |
| default_writer_properties(), coerce_millis)); |
| ASSERT_RAISES(Invalid, WriteTable(*t4, ::arrow::default_memory_pool(), sink, 10, |
| default_writer_properties(), coerce_millis)); |
| |
| // OK to write micros to micros |
| auto coerce_micros = |
| (ArrowWriterProperties::Builder().coerce_timestamps(TimeUnit::MICRO)->build()); |
| ASSERT_OK_NO_THROW(WriteTable(*t3, ::arrow::default_memory_pool(), sink, 10, |
| default_writer_properties(), coerce_micros)); |
| |
| // Loss of precision |
| ASSERT_RAISES(Invalid, WriteTable(*t4, ::arrow::default_memory_pool(), sink, 10, |
| default_writer_properties(), coerce_micros)); |
| } |
| |
| TEST(TestArrowReadWrite, ConvertedDateTimeTypes) { |
| using ::arrow::ArrayFromVector; |
| |
| std::vector<bool> is_valid = {true, true, true, false, true, true}; |
| |
| auto f0 = field("f0", ::arrow::date64()); |
| auto f1 = field("f1", ::arrow::time32(TimeUnit::SECOND)); |
| std::shared_ptr<::arrow::Schema> schema(new ::arrow::Schema({f0, f1})); |
| |
| std::vector<int64_t> a0_values = {1489190400000, 1489276800000, 1489363200000, |
| 1489449600000, 1489536000000, 1489622400000}; |
| std::vector<int32_t> a1_values = {0, 1, 2, 3, 4, 5}; |
| |
| std::shared_ptr<Array> a0, a1, x0, x1; |
| ArrayFromVector<::arrow::Date64Type, int64_t>(f0->type(), is_valid, a0_values, &a0); |
| ArrayFromVector<::arrow::Time32Type, int32_t>(f1->type(), is_valid, a1_values, &a1); |
| |
| std::vector<std::shared_ptr<::arrow::Column>> columns = { |
| std::make_shared<Column>("f0", a0), std::make_shared<Column>("f1", a1)}; |
| auto table = std::make_shared<::arrow::Table>(schema, columns); |
| |
| // Expected schema and values |
| auto e0 = field("f0", ::arrow::date32()); |
| auto e1 = field("f1", ::arrow::time32(TimeUnit::MILLI)); |
| std::shared_ptr<::arrow::Schema> ex_schema(new ::arrow::Schema({e0, e1})); |
| |
| std::vector<int32_t> x0_values = {17236, 17237, 17238, 17239, 17240, 17241}; |
| std::vector<int32_t> x1_values = {0, 1000, 2000, 3000, 4000, 5000}; |
| ArrayFromVector<::arrow::Date32Type, int32_t>(e0->type(), is_valid, x0_values, &x0); |
| ArrayFromVector<::arrow::Time32Type, int32_t>(e1->type(), is_valid, x1_values, &x1); |
| |
| std::vector<std::shared_ptr<::arrow::Column>> ex_columns = { |
| std::make_shared<Column>("f0", x0), std::make_shared<Column>("f1", x1)}; |
| auto ex_table = std::make_shared<::arrow::Table>(ex_schema, ex_columns); |
| |
| std::shared_ptr<Table> result; |
| DoSimpleRoundtrip(table, 1, table->num_rows(), {}, &result); |
| |
| AssertTablesEqual(*ex_table, *result); |
| } |
| |
| void MakeDoubleTable(int num_columns, int num_rows, int nchunks, |
| std::shared_ptr<Table>* out) { |
| std::shared_ptr<::arrow::Column> column; |
| std::vector<std::shared_ptr<::arrow::Column>> columns(num_columns); |
| std::vector<std::shared_ptr<::arrow::Field>> fields(num_columns); |
| |
| for (int i = 0; i < num_columns; ++i) { |
| std::vector<std::shared_ptr<Array>> arrays; |
| std::shared_ptr<Array> values; |
| ASSERT_OK(NullableArray<::arrow::DoubleType>(num_rows, num_rows / 10, |
| static_cast<uint32_t>(i), &values)); |
| std::stringstream ss; |
| ss << "col" << i; |
| |
| for (int j = 0; j < nchunks; ++j) { |
| arrays.push_back(values); |
| } |
| column = MakeColumn(ss.str(), arrays, true); |
| |
| columns[i] = column; |
| fields[i] = column->field(); |
| } |
| auto schema = std::make_shared<::arrow::Schema>(fields); |
| *out = std::make_shared<Table>(schema, columns); |
| } |
| |
| TEST(TestArrowReadWrite, MultithreadedRead) { |
| const int num_columns = 20; |
| const int num_rows = 1000; |
| const int num_threads = 4; |
| |
| std::shared_ptr<Table> table; |
| MakeDoubleTable(num_columns, num_rows, 1, &table); |
| |
| std::shared_ptr<Table> result; |
| DoSimpleRoundtrip(table, num_threads, table->num_rows(), {}, &result); |
| |
| AssertTablesEqual(*table, *result); |
| } |
| |
| TEST(TestArrowReadWrite, ReadSingleRowGroup) { |
| const int num_columns = 20; |
| const int num_rows = 1000; |
| |
| std::shared_ptr<Table> table; |
| MakeDoubleTable(num_columns, num_rows, 1, &table); |
| |
| std::shared_ptr<Buffer> buffer; |
| WriteTableToBuffer(table, 1, num_rows / 2, default_arrow_writer_properties(), &buffer); |
| |
| std::unique_ptr<FileReader> reader; |
| ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer), |
| ::arrow::default_memory_pool(), |
| ::parquet::default_reader_properties(), nullptr, &reader)); |
| |
| ASSERT_EQ(2, reader->num_row_groups()); |
| |
| std::shared_ptr<Table> r1, r2; |
| // Read everything |
| ASSERT_OK_NO_THROW(reader->ReadRowGroup(0, &r1)); |
| ASSERT_OK_NO_THROW(reader->ReadRowGroup(1, &r2)); |
| |
| std::shared_ptr<Table> concatenated; |
| ASSERT_OK(ConcatenateTables({r1, r2}, &concatenated)); |
| |
| ASSERT_TRUE(table->Equals(*concatenated)); |
| } |
| |
| TEST(TestArrowReadWrite, ScanContents) { |
| const int num_columns = 20; |
| const int num_rows = 1000; |
| |
| std::shared_ptr<Table> table; |
| MakeDoubleTable(num_columns, num_rows, 1, &table); |
| |
| std::shared_ptr<Buffer> buffer; |
| WriteTableToBuffer(table, 1, num_rows / 2, default_arrow_writer_properties(), &buffer); |
| |
| std::unique_ptr<FileReader> reader; |
| ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer), |
| ::arrow::default_memory_pool(), |
| ::parquet::default_reader_properties(), nullptr, &reader)); |
| |
| int64_t num_rows_returned = 0; |
| ASSERT_OK_NO_THROW(reader->ScanContents({}, 256, &num_rows_returned)); |
| ASSERT_EQ(num_rows, num_rows_returned); |
| |
| ASSERT_OK_NO_THROW(reader->ScanContents({0, 1, 2}, 256, &num_rows_returned)); |
| ASSERT_EQ(num_rows, num_rows_returned); |
| } |
| |
| TEST(TestArrowReadWrite, ReadColumnSubset) { |
| const int num_columns = 20; |
| const int num_rows = 1000; |
| const int num_threads = 4; |
| |
| std::shared_ptr<Table> table; |
| MakeDoubleTable(num_columns, num_rows, 1, &table); |
| |
| std::shared_ptr<Table> result; |
| std::vector<int> column_subset = {0, 4, 8, 10}; |
| DoSimpleRoundtrip(table, num_threads, table->num_rows(), column_subset, &result); |
| |
| std::vector<std::shared_ptr<::arrow::Column>> ex_columns; |
| std::vector<std::shared_ptr<::arrow::Field>> ex_fields; |
| for (int i : column_subset) { |
| ex_columns.push_back(table->column(i)); |
| ex_fields.push_back(table->column(i)->field()); |
| } |
| |
| auto ex_schema = std::make_shared<::arrow::Schema>(ex_fields); |
| Table expected(ex_schema, ex_columns); |
| AssertTablesEqual(expected, *result); |
| } |
| |
| void MakeListTable(int num_rows, std::shared_ptr<Table>* out) { |
| ::arrow::Int32Builder offset_builder; |
| |
| std::vector<int32_t> length_draws; |
| randint(num_rows, 0, 100, &length_draws); |
| |
| std::vector<int32_t> offset_values; |
| |
| // Make sure some of them are length 0 |
| int32_t total_elements = 0; |
| for (size_t i = 0; i < length_draws.size(); ++i) { |
| if (length_draws[i] < 10) { |
| length_draws[i] = 0; |
| } |
| offset_values.push_back(total_elements); |
| total_elements += length_draws[i]; |
| } |
| offset_values.push_back(total_elements); |
| |
| std::vector<int8_t> value_draws; |
| randint<int8_t>(total_elements, 0, 100, &value_draws); |
| |
| std::vector<bool> is_valid; |
| random_is_valid(total_elements, 0.1, &is_valid); |
| |
| std::shared_ptr<Array> values, offsets; |
| ::arrow::ArrayFromVector<::arrow::Int8Type, int8_t>(::arrow::int8(), is_valid, |
| value_draws, &values); |
| ::arrow::ArrayFromVector<::arrow::Int32Type, int32_t>(offset_values, &offsets); |
| |
| std::shared_ptr<Array> list_array; |
| ASSERT_OK(::arrow::ListArray::FromArrays(*offsets, *values, default_memory_pool(), |
| &list_array)); |
| |
| auto f1 = ::arrow::field("a", ::arrow::list(::arrow::int8())); |
| auto schema = ::arrow::schema({f1}); |
| std::vector<std::shared_ptr<Array>> arrays = {list_array}; |
| *out = std::make_shared<Table>(schema, arrays); |
| } |
| |
| TEST(TestArrowReadWrite, ListLargeRecords) { |
| const int num_rows = 50; |
| |
| std::shared_ptr<Table> table; |
| MakeListTable(num_rows, &table); |
| |
| std::shared_ptr<Buffer> buffer; |
| WriteTableToBuffer(table, 1, 100, default_arrow_writer_properties(), &buffer); |
| |
| std::unique_ptr<FileReader> reader; |
| ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer), |
| ::arrow::default_memory_pool(), |
| ::parquet::default_reader_properties(), nullptr, &reader)); |
| |
| // Read everything |
| std::shared_ptr<Table> result; |
| ASSERT_OK_NO_THROW(reader->ReadTable(&result)); |
| AssertTablesEqual(*table, *result); |
| |
| ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer), |
| ::arrow::default_memory_pool(), |
| ::parquet::default_reader_properties(), nullptr, &reader)); |
| |
| std::unique_ptr<ColumnReader> col_reader; |
| ASSERT_OK(reader->GetColumn(0, &col_reader)); |
| |
| auto expected = table->column(0)->data()->chunk(0); |
| |
| std::vector<std::shared_ptr<Array>> pieces; |
| for (int i = 0; i < num_rows; ++i) { |
| std::shared_ptr<Array> piece; |
| ASSERT_OK(col_reader->NextBatch(1, &piece)); |
| ASSERT_EQ(1, piece->length()); |
| pieces.push_back(piece); |
| } |
| auto chunked = std::make_shared<::arrow::ChunkedArray>(pieces); |
| |
| auto chunked_col = |
| std::make_shared<::arrow::Column>(table->schema()->field(0), chunked); |
| std::vector<std::shared_ptr<::arrow::Column>> columns = {chunked_col}; |
| auto chunked_table = std::make_shared<Table>(table->schema(), columns); |
| |
| ASSERT_TRUE(table->Equals(*chunked_table)); |
| } |
| |
| TEST(TestArrowWrite, CheckChunkSize) { |
| const int num_columns = 2; |
| const int num_rows = 128; |
| const int64_t chunk_size = 0; // note the chunk_size is 0 |
| std::shared_ptr<Table> table; |
| MakeDoubleTable(num_columns, num_rows, 1, &table); |
| |
| auto sink = std::make_shared<InMemoryOutputStream>(); |
| |
| ASSERT_RAISES(Invalid, |
| WriteTable(*table, ::arrow::default_memory_pool(), sink, chunk_size)); |
| } |
| |
| class TestNestedSchemaRead : public ::testing::TestWithParam<Repetition::type> { |
| protected: |
| // make it *3 to make it easily divisible by 3 |
| const int NUM_SIMPLE_TEST_ROWS = SMALL_SIZE * 3; |
| std::shared_ptr<::arrow::Int32Array> values_array_ = nullptr; |
| |
| void InitReader() { |
| std::shared_ptr<Buffer> buffer = nested_parquet_->GetBuffer(); |
| ASSERT_OK_NO_THROW( |
| OpenFile(std::make_shared<BufferReader>(buffer), ::arrow::default_memory_pool(), |
| ::parquet::default_reader_properties(), nullptr, &reader_)); |
| } |
| |
| void InitNewParquetFile(const std::shared_ptr<GroupNode>& schema, int num_rows) { |
| nested_parquet_ = std::make_shared<InMemoryOutputStream>(); |
| |
| writer_ = parquet::ParquetFileWriter::Open(nested_parquet_, schema, |
| default_writer_properties()); |
| row_group_writer_ = writer_->AppendRowGroup(); |
| } |
| |
| void FinalizeParquetFile() { |
| row_group_writer_->Close(); |
| writer_->Close(); |
| } |
| |
| void MakeValues(int num_rows) { |
| std::shared_ptr<Array> arr; |
| ASSERT_OK(NullableArray<::arrow::Int32Type>(num_rows, 0, kDefaultSeed, &arr)); |
| values_array_ = std::dynamic_pointer_cast<::arrow::Int32Array>(arr); |
| } |
| |
| void WriteColumnData(size_t num_rows, int16_t* def_levels, int16_t* rep_levels, |
| int32_t* values) { |
| auto typed_writer = |
| static_cast<TypedColumnWriter<Int32Type>*>(row_group_writer_->NextColumn()); |
| typed_writer->WriteBatch(num_rows, def_levels, rep_levels, values); |
| } |
| |
| void ValidateArray(const Array& array, size_t expected_nulls) { |
| ASSERT_EQ(array.length(), values_array_->length()); |
| ASSERT_EQ(array.null_count(), expected_nulls); |
| // Also independently count the nulls |
| auto local_null_count = 0; |
| for (int i = 0; i < array.length(); i++) { |
| if (array.IsNull(i)) { |
| local_null_count++; |
| } |
| } |
| ASSERT_EQ(local_null_count, expected_nulls); |
| } |
| |
| void ValidateColumnArray(const ::arrow::Int32Array& array, size_t expected_nulls) { |
| ValidateArray(array, expected_nulls); |
| int j = 0; |
| for (int i = 0; i < values_array_->length(); i++) { |
| if (array.IsNull(i)) { |
| continue; |
| } |
| ASSERT_EQ(array.Value(i), values_array_->Value(j)); |
| j++; |
| } |
| } |
| |
| void ValidateTableArrayTypes(const Table& table) { |
| for (int i = 0; i < table.num_columns(); i++) { |
| const std::shared_ptr<::arrow::Field> schema_field = table.schema()->field(i); |
| const std::shared_ptr<Column> column = table.column(i); |
| // Compare with the column field |
| ASSERT_TRUE(schema_field->Equals(column->field())); |
| // Compare with the array type |
| ASSERT_TRUE(schema_field->type()->Equals(column->data()->chunk(0)->type())); |
| } |
| } |
| |
| // A parquet with a simple nested schema |
| void CreateSimpleNestedParquet(Repetition::type struct_repetition) { |
| std::vector<NodePtr> parquet_fields; |
| // TODO(itaiin): We are using parquet low-level file api to create the nested parquet |
| // this needs to change when a nested writes are implemented |
| |
| // create the schema: |
| // <struct_repetition> group group1 { |
| // required int32 leaf1; |
| // optional int32 leaf2; |
| // } |
| // required int32 leaf3; |
| |
| parquet_fields.push_back(GroupNode::Make( |
| "group1", struct_repetition, |
| {PrimitiveNode::Make("leaf1", Repetition::REQUIRED, ParquetType::INT32), |
| PrimitiveNode::Make("leaf2", Repetition::OPTIONAL, ParquetType::INT32)})); |
| parquet_fields.push_back( |
| PrimitiveNode::Make("leaf3", Repetition::REQUIRED, ParquetType::INT32)); |
| |
| auto schema_node = GroupNode::Make("schema", Repetition::REQUIRED, parquet_fields); |
| |
| // Create definition levels for the different columns that contain interleaved |
| // nulls and values at all nesting levels |
| |
| // definition levels for optional fields |
| std::vector<int16_t> leaf1_def_levels(NUM_SIMPLE_TEST_ROWS); |
| std::vector<int16_t> leaf2_def_levels(NUM_SIMPLE_TEST_ROWS); |
| std::vector<int16_t> leaf3_def_levels(NUM_SIMPLE_TEST_ROWS); |
| for (int i = 0; i < NUM_SIMPLE_TEST_ROWS; i++) { |
| // leaf1 is required within the optional group1, so it is only null |
| // when the group is null |
| leaf1_def_levels[i] = (i % 3 == 0) ? 0 : 1; |
| // leaf2 is optional, can be null in the primitive (def-level 1) or |
| // struct level (def-level 0) |
| leaf2_def_levels[i] = i % 3; |
| // leaf3 is required |
| leaf3_def_levels[i] = 0; |
| } |
| |
| std::vector<int16_t> rep_levels(NUM_SIMPLE_TEST_ROWS, 0); |
| |
| // Produce values for the columns |
| MakeValues(NUM_SIMPLE_TEST_ROWS); |
| int32_t* values = reinterpret_cast<int32_t*>(values_array_->values()->mutable_data()); |
| |
| // Create the actual parquet file |
| InitNewParquetFile(std::static_pointer_cast<GroupNode>(schema_node), |
| NUM_SIMPLE_TEST_ROWS); |
| |
| // leaf1 column |
| WriteColumnData(NUM_SIMPLE_TEST_ROWS, leaf1_def_levels.data(), rep_levels.data(), |
| values); |
| // leaf2 column |
| WriteColumnData(NUM_SIMPLE_TEST_ROWS, leaf2_def_levels.data(), rep_levels.data(), |
| values); |
| // leaf3 column |
| WriteColumnData(NUM_SIMPLE_TEST_ROWS, leaf3_def_levels.data(), rep_levels.data(), |
| values); |
| |
| FinalizeParquetFile(); |
| InitReader(); |
| } |
| |
| NodePtr CreateSingleTypedNestedGroup(int index, int depth, int num_children, |
| Repetition::type node_repetition, |
| ParquetType::type leaf_type) { |
| std::vector<NodePtr> children; |
| |
| for (int i = 0; i < num_children; i++) { |
| if (depth <= 1) { |
| children.push_back(PrimitiveNode::Make("leaf", node_repetition, leaf_type)); |
| } else { |
| children.push_back(CreateSingleTypedNestedGroup(i, depth - 1, num_children, |
| node_repetition, leaf_type)); |
| } |
| } |
| |
| std::stringstream ss; |
| ss << "group-" << depth << "-" << index; |
| return NodePtr(GroupNode::Make(ss.str(), node_repetition, children)); |
| } |
| |
| // A deeply nested schema |
| void CreateMultiLevelNestedParquet(int num_trees, int tree_depth, int num_children, |
| int num_rows, Repetition::type node_repetition) { |
| // Create the schema |
| std::vector<NodePtr> parquet_fields; |
| for (int i = 0; i < num_trees; i++) { |
| parquet_fields.push_back(CreateSingleTypedNestedGroup( |
| i, tree_depth, num_children, node_repetition, ParquetType::INT32)); |
| } |
| auto schema_node = GroupNode::Make("schema", Repetition::REQUIRED, parquet_fields); |
| |
| int num_columns = num_trees * static_cast<int>((std::pow(num_children, tree_depth))); |
| |
| std::vector<int16_t> def_levels; |
| std::vector<int16_t> rep_levels; |
| |
| int num_levels = 0; |
| while (num_levels < num_rows) { |
| if (node_repetition == Repetition::REQUIRED) { |
| def_levels.push_back(0); // all are required |
| } else { |
| int16_t level = static_cast<int16_t>(num_levels % (tree_depth + 2)); |
| def_levels.push_back(level); // all are optional |
| } |
| rep_levels.push_back(0); // none is repeated |
| ++num_levels; |
| } |
| |
| // Produce values for the columns |
| MakeValues(num_rows); |
| int32_t* values = reinterpret_cast<int32_t*>(values_array_->values()->mutable_data()); |
| |
| // Create the actual parquet file |
| InitNewParquetFile(std::static_pointer_cast<GroupNode>(schema_node), num_rows); |
| |
| for (int i = 0; i < num_columns; i++) { |
| WriteColumnData(num_rows, def_levels.data(), rep_levels.data(), values); |
| } |
| FinalizeParquetFile(); |
| InitReader(); |
| } |
| |
| class DeepParquetTestVisitor : public ArrayVisitor { |
| public: |
| DeepParquetTestVisitor(Repetition::type node_repetition, |
| std::shared_ptr<::arrow::Int32Array> expected) |
| : node_repetition_(node_repetition), expected_(expected) {} |
| |
| Status Validate(std::shared_ptr<Array> tree) { return tree->Accept(this); } |
| |
| virtual Status Visit(const ::arrow::Int32Array& array) { |
| if (node_repetition_ == Repetition::REQUIRED) { |
| if (!array.Equals(expected_)) { |
| return Status::Invalid("leaf array data mismatch"); |
| } |
| } else if (node_repetition_ == Repetition::OPTIONAL) { |
| if (array.length() != expected_->length()) { |
| return Status::Invalid("Bad leaf array length"); |
| } |
| // expect only 1 value every `depth` row |
| if (array.null_count() != SMALL_SIZE) { |
| return Status::Invalid("Unexpected null count"); |
| } |
| } else { |
| return Status::NotImplemented("Unsupported repetition"); |
| } |
| return Status::OK(); |
| } |
| |
| virtual Status Visit(const ::arrow::StructArray& array) { |
| for (int32_t i = 0; i < array.num_fields(); ++i) { |
| auto child = array.field(i); |
| if (node_repetition_ == Repetition::REQUIRED) { |
| RETURN_NOT_OK(child->Accept(this)); |
| } else if (node_repetition_ == Repetition::OPTIONAL) { |
| // Null count Must be a multiple of SMALL_SIZE |
| if (array.null_count() % SMALL_SIZE != 0) { |
| return Status::Invalid("Unexpected struct null count"); |
| } |
| } else { |
| return Status::NotImplemented("Unsupported repetition"); |
| } |
| } |
| return Status::OK(); |
| } |
| |
| private: |
| Repetition::type node_repetition_; |
| std::shared_ptr<::arrow::Int32Array> expected_; |
| }; |
| |
| std::shared_ptr<InMemoryOutputStream> nested_parquet_; |
| std::unique_ptr<FileReader> reader_; |
| std::unique_ptr<ParquetFileWriter> writer_; |
| RowGroupWriter* row_group_writer_; |
| }; |
| |
| TEST_F(TestNestedSchemaRead, ReadIntoTableFull) { |
| CreateSimpleNestedParquet(Repetition::OPTIONAL); |
| |
| std::shared_ptr<Table> table; |
| ASSERT_OK_NO_THROW(reader_->ReadTable(&table)); |
| ASSERT_EQ(table->num_rows(), NUM_SIMPLE_TEST_ROWS); |
| ASSERT_EQ(table->num_columns(), 2); |
| ASSERT_EQ(table->schema()->field(0)->type()->num_children(), 2); |
| ValidateTableArrayTypes(*table); |
| |
| auto struct_field_array = |
| std::static_pointer_cast<::arrow::StructArray>(table->column(0)->data()->chunk(0)); |
| auto leaf1_array = |
| std::static_pointer_cast<::arrow::Int32Array>(struct_field_array->field(0)); |
| auto leaf2_array = |
| std::static_pointer_cast<::arrow::Int32Array>(struct_field_array->field(1)); |
| auto leaf3_array = |
| std::static_pointer_cast<::arrow::Int32Array>(table->column(1)->data()->chunk(0)); |
| |
| // validate struct and leaf arrays |
| |
| // validate struct array |
| ValidateArray(*struct_field_array, NUM_SIMPLE_TEST_ROWS / 3); |
| // validate leaf1 |
| ValidateColumnArray(*leaf1_array, NUM_SIMPLE_TEST_ROWS / 3); |
| // validate leaf2 |
| ValidateColumnArray(*leaf2_array, NUM_SIMPLE_TEST_ROWS * 2 / 3); |
| // validate leaf3 |
| ValidateColumnArray(*leaf3_array, 0); |
| } |
| |
| TEST_F(TestNestedSchemaRead, ReadTablePartial) { |
| CreateSimpleNestedParquet(Repetition::OPTIONAL); |
| std::shared_ptr<Table> table; |
| |
| // columns: {group1.leaf1, leaf3} |
| ASSERT_OK_NO_THROW(reader_->ReadTable({0, 2}, &table)); |
| ASSERT_EQ(table->num_rows(), NUM_SIMPLE_TEST_ROWS); |
| ASSERT_EQ(table->num_columns(), 2); |
| ASSERT_EQ(table->schema()->field(0)->name(), "group1"); |
| ASSERT_EQ(table->schema()->field(1)->name(), "leaf3"); |
| ASSERT_EQ(table->schema()->field(0)->type()->num_children(), 1); |
| ValidateTableArrayTypes(*table); |
| |
| // columns: {group1.leaf1, group1.leaf2} |
| ASSERT_OK_NO_THROW(reader_->ReadTable({0, 1}, &table)); |
| ASSERT_EQ(table->num_rows(), NUM_SIMPLE_TEST_ROWS); |
| ASSERT_EQ(table->num_columns(), 1); |
| ASSERT_EQ(table->schema()->field(0)->name(), "group1"); |
| ASSERT_EQ(table->schema()->field(0)->type()->num_children(), 2); |
| ValidateTableArrayTypes(*table); |
| |
| // columns: {leaf3} |
| ASSERT_OK_NO_THROW(reader_->ReadTable({2}, &table)); |
| ASSERT_EQ(table->num_rows(), NUM_SIMPLE_TEST_ROWS); |
| ASSERT_EQ(table->num_columns(), 1); |
| ASSERT_EQ(table->schema()->field(0)->name(), "leaf3"); |
| ASSERT_EQ(table->schema()->field(0)->type()->num_children(), 0); |
| ValidateTableArrayTypes(*table); |
| |
| // Test with different ordering |
| ASSERT_OK_NO_THROW(reader_->ReadTable({2, 0}, &table)); |
| ASSERT_EQ(table->num_rows(), NUM_SIMPLE_TEST_ROWS); |
| ASSERT_EQ(table->num_columns(), 2); |
| ASSERT_EQ(table->schema()->field(0)->name(), "leaf3"); |
| ASSERT_EQ(table->schema()->field(1)->name(), "group1"); |
| ASSERT_EQ(table->schema()->field(1)->type()->num_children(), 1); |
| ValidateTableArrayTypes(*table); |
| } |
| |
| TEST_F(TestNestedSchemaRead, StructAndListTogetherUnsupported) { |
| CreateSimpleNestedParquet(Repetition::REPEATED); |
| std::shared_ptr<Table> table; |
| ASSERT_RAISES(NotImplemented, reader_->ReadTable(&table)); |
| } |
| |
| TEST_P(TestNestedSchemaRead, DeepNestedSchemaRead) { |
| const int num_trees = 10; |
| const int depth = 5; |
| const int num_children = 3; |
| int num_rows = SMALL_SIZE * (depth + 2); |
| CreateMultiLevelNestedParquet(num_trees, depth, num_children, num_rows, GetParam()); |
| std::shared_ptr<Table> table; |
| ASSERT_OK_NO_THROW(reader_->ReadTable(&table)); |
| ASSERT_EQ(table->num_columns(), num_trees); |
| ASSERT_EQ(table->num_rows(), num_rows); |
| |
| DeepParquetTestVisitor visitor(GetParam(), values_array_); |
| for (int i = 0; i < table->num_columns(); i++) { |
| auto tree = table->column(i)->data()->chunk(0); |
| ASSERT_OK_NO_THROW(visitor.Validate(tree)); |
| } |
| } |
| |
| INSTANTIATE_TEST_CASE_P(Repetition_type, TestNestedSchemaRead, |
| ::testing::Values(Repetition::REQUIRED, Repetition::OPTIONAL)); |
| |
| TEST(TestImpalaConversion, NanosecondToImpala) { |
| // June 20, 2017 16:32:56 and 123456789 nanoseconds |
| int64_t nanoseconds = INT64_C(1497976376123456789); |
| Int96 expected = {{UINT32_C(632093973), UINT32_C(13871), UINT32_C(2457925)}}; |
| Int96 calculated; |
| internal::NanosecondsToImpalaTimestamp(nanoseconds, &calculated); |
| ASSERT_EQ(expected, calculated); |
| } |
| |
| TEST(TestArrowReaderAdHoc, Int96BadMemoryAccess) { |
| // PARQUET-995 |
| const char* data_dir = std::getenv("PARQUET_TEST_DATA"); |
| std::string dir_string(data_dir); |
| std::stringstream ss; |
| ss << dir_string << "/" |
| << "alltypes_plain.parquet"; |
| auto path = ss.str(); |
| |
| auto pool = ::arrow::default_memory_pool(); |
| |
| std::unique_ptr<FileReader> arrow_reader; |
| ASSERT_NO_THROW( |
| arrow_reader.reset(new FileReader(pool, ParquetFileReader::OpenFile(path, false)))); |
| std::shared_ptr<::arrow::Table> table; |
| ASSERT_OK_NO_THROW(arrow_reader->ReadTable(&table)); |
| } |
| |
| } // namespace arrow |
| |
| } // namespace parquet |