ARROW-91: Basic Parquet read support

cpp/src/arrow/parquet/CMakeLists.txt

@@ -19,6 +19,7 @@
 # arrow_parquet : Arrow <-> Parquet adapter
 
 set(PARQUET_SRCS
+ reader.cc
  schema.cc
 )
 
@@ -36,6 +37,9 @@ SET_TARGET_PROPERTIES(arrow_parquet PROPERTIES LINKER_LANGUAGE CXX)
 ADD_ARROW_TEST(parquet-schema-test)
 ARROW_TEST_LINK_LIBRARIES(parquet-schema-test arrow_parquet)
 
+ADD_ARROW_TEST(parquet-reader-test)
+ARROW_TEST_LINK_LIBRARIES(parquet-reader-test arrow_parquet)
+
 # Headers: top level
 install(FILES
  DESTINATION include/arrow/parquet)

cpp/src/arrow/parquet/parquet-reader-test.cc

@@ -0,0 +1,120 @@
+// 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.
+
+#include "gtest/gtest.h"
+
+#include "arrow/test-util.h"
+#include "arrow/parquet/reader.h"
+#include "arrow/types/primitive.h"
+#include "arrow/util/memory-pool.h"
+#include "arrow/util/status.h"
+
+#include "parquet/api/schema.h"
+#include "parquet/column/writer.h"
+#include "parquet/file/reader.h"
+#include "parquet/file/writer.h"
+#include "parquet/util/input.h"
+#include "parquet/util/output.h"
+
+using ParquetBuffer = parquet::Buffer;
+using parquet::BufferReader;
+using parquet::InMemoryOutputStream;
+using parquet::Int64Writer;
+using parquet::ParquetFileReader;
+using parquet::ParquetFileWriter;
+using parquet::RandomAccessSource;
+using parquet::Repetition;
+using parquet::SchemaDescriptor;
+using ParquetType = parquet::Type;
+using parquet::schema::GroupNode;
+using parquet::schema::NodePtr;
+using parquet::schema::PrimitiveNode;
+
+namespace arrow {
+
+namespace parquet {
+
+class TestReadParquet : public ::testing::Test {
+ public:
+ virtual void SetUp() {}
+
+ std::shared_ptr<GroupNode> Int64Schema() {
+ auto pnode = PrimitiveNode::Make("int64", Repetition::REQUIRED, ParquetType::INT64);
+ NodePtr node_ =
+ GroupNode::Make("schema", Repetition::REQUIRED, std::vector<NodePtr>({pnode}));
+ return std::static_pointer_cast<GroupNode>(node_);
+ }
+
+ std::unique_ptr<ParquetFileReader> Int64File(
+ std::vector<int64_t>& values, int num_chunks) {
+ std::shared_ptr<GroupNode> schema = Int64Schema();
+ std::shared_ptr<InMemoryOutputStream> sink(new InMemoryOutputStream());
+ auto file_writer = ParquetFileWriter::Open(sink, 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<Int64Writer*>(row_group_writer->NextColumn());
+ int64_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();
+
+ std::shared_ptr<ParquetBuffer> buffer = sink->GetBuffer();
+ std::unique_ptr<RandomAccessSource> source(new BufferReader(buffer));
+ return ParquetFileReader::Open(std::move(source));
+ }
+
+ private:
+};
+
+TEST_F(TestReadParquet, SingleColumnInt64) {
+ std::vector<int64_t> values(100, 128);
+ std::unique_ptr<ParquetFileReader> file_reader = Int64File(values, 1);
+ arrow::parquet::FileReader reader(default_memory_pool(), std::move(file_reader));
+ std::unique_ptr<arrow::parquet::FlatColumnReader> column_reader;
+ ASSERT_OK(reader.GetFlatColumn(0, &column_reader));
+ ASSERT_NE(nullptr, column_reader.get());
+ std::shared_ptr<Array> out;
+ ASSERT_OK(column_reader->NextBatch(100, &out));
+ ASSERT_NE(nullptr, out.get());
+ 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]);
+ }
+}
+
+TEST_F(TestReadParquet, SingleColumnInt64Chunked) {
+ std::vector<int64_t> values(100, 128);
+ std::unique_ptr<ParquetFileReader> file_reader = Int64File(values, 4);
+ arrow::parquet::FileReader reader(default_memory_pool(), std::move(file_reader));
+ std::unique_ptr<arrow::parquet::FlatColumnReader> column_reader;
+ ASSERT_OK(reader.GetFlatColumn(0, &column_reader));
+ ASSERT_NE(nullptr, column_reader.get());
+ std::shared_ptr<Array> out;
+ ASSERT_OK(column_reader->NextBatch(100, &out));
+ ASSERT_NE(nullptr, out.get());
+ 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]);
+ }
+}
+
+} // namespace parquet
+
+} // namespace arrow

cpp/src/arrow/parquet/reader.cc

@@ -0,0 +1,157 @@
+// 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.
+
+#include "arrow/parquet/reader.h"
+
+#include <queue>
+
+#include "arrow/parquet/schema.h"
+#include "arrow/schema.h"
+#include "arrow/types/primitive.h"
+#include "arrow/util/status.h"
+
+using parquet::ColumnReader;
+using parquet::TypedColumnReader;
+
+namespace arrow {
+namespace parquet {
+
+class FileReader::Impl {
+ public:
+ Impl(MemoryPool* pool, std::unique_ptr<::parquet::ParquetFileReader> reader);
+ virtual ~Impl() {}
+
+ Status GetFlatColumn(int i, std::unique_ptr<FlatColumnReader>* out);
+
+ private:
+ MemoryPool* pool_;
+ std::unique_ptr<::parquet::ParquetFileReader> reader_;
+};
+
+class FlatColumnReader::Impl {
+ public:
+ Impl(MemoryPool* pool, const ::parquet::ColumnDescriptor* descr,
+ std::queue<std::shared_ptr<ColumnReader>>&& column_readers);
+ virtual ~Impl() {}
+
+ Status NextBatch(int batch_size, std::shared_ptr<Array>* out);
+ template <typename ArrowType, typename ParquetType, typename CType>
+ Status TypedReadBatch(int batch_size, std::shared_ptr<Array>* out);
+
+ private:
+ MemoryPool* pool_;
+ const ::parquet::ColumnDescriptor* descr_;
+ std::queue<std::shared_ptr<ColumnReader>> column_readers_;
+ std::shared_ptr<Field> field_;
+};
+
+FileReader::Impl::Impl(
+ MemoryPool* pool, std::unique_ptr<::parquet::ParquetFileReader> reader)
+ : pool_(pool), reader_(std::move(reader)) {}
+
+Status FileReader::Impl::GetFlatColumn(int i, std::unique_ptr<FlatColumnReader>* out) {
+ std::queue<std::shared_ptr<ColumnReader>> column_readers;
+ for (int rg = 0; rg < reader_->num_row_groups(); rg++) {
+ column_readers.push(reader_->RowGroup(rg)->Column(i));
+ }
+ std::unique_ptr<FlatColumnReader::Impl> impl(new FlatColumnReader::Impl(
+ pool_, reader_->descr()->Column(i), std::move(column_readers)));
+ *out = std::unique_ptr<FlatColumnReader>(new FlatColumnReader(std::move(impl)));
+ return Status::OK();
+}
+
+FileReader::FileReader(
+ MemoryPool* pool, std::unique_ptr<::parquet::ParquetFileReader> reader)
+ : impl_(new FileReader::Impl(pool, std::move(reader))) {}
+
+FileReader::~FileReader() {}
+
+Status FileReader::GetFlatColumn(int i, std::unique_ptr<FlatColumnReader>* out) {
+ return impl_->GetFlatColumn(i, out);
+}
+
+Status FileReader::ReadFlatColumn(int i, std::shared_ptr<Array>* out) {
+ return Status::OK();
+}
+
+FlatColumnReader::Impl::Impl(MemoryPool* pool, const ::parquet::ColumnDescriptor* descr,
+ std::queue<std::shared_ptr<ColumnReader>>&& column_readers)
+ : pool_(pool), descr_(descr), column_readers_(column_readers) {
+ NodeToField(descr_->schema_node(), &field_);
+}
+
+template <typename ArrowType, typename ParquetType, typename CType>
+Status FlatColumnReader::Impl::TypedReadBatch(
+ int batch_size, std::shared_ptr<Array>* out) {
+ int values_to_read = batch_size;
+ NumericBuilder<ArrowType> builder(pool_, field_->type);
+ while ((values_to_read > 0) && (column_readers_.size() > 0)) {
+ // TODO: This is a lot malloc-thresing and not using the memory pool.
+ std::vector<CType> values(values_to_read);
+ std::vector<int16_t> def_levels(values_to_read);
+ auto reader =
+ dynamic_cast<TypedColumnReader<ParquetType>*>(column_readers_.front().get());
+ int64_t values_read;
+ values_to_read -= reader->ReadBatch(
+ values_to_read, def_levels.data(), nullptr, values.data(), &values_read);
+ if (descr_->max_definition_level() == 0) {
+ builder.Append(values.data(), values_read);
+ } else {
+ return Status::NotImplemented("no support for definition levels yet");
+ }
+ if (!column_readers_.front()->HasNext()) { column_readers_.pop(); }
+ }
+ *out = builder.Finish();
+ return Status::OK();
+}
+
+#define TYPED_BATCH_CASE(ENUM, ArrowType, ParquetType, CType) \
+ case Type::ENUM: \
+ return TypedReadBatch<ArrowType, ParquetType, CType>(batch_size, out); \
+ break;
+
+Status FlatColumnReader::Impl::NextBatch(int batch_size, std::shared_ptr<Array>* out) {
+ if (column_readers_.size() == 0) {
+ // Exhausted all readers.
+ *out = std::shared_ptr<Array>(nullptr);
+ }
+
+ if (descr_->max_repetition_level() > 0) {
+ return Status::NotImplemented("no support for repetition yet");
+ }
+
+ *out = std::shared_ptr<Array>(nullptr);
+ switch (field_->type->type) {
+ TYPED_BATCH_CASE(INT32, Int32Type, ::parquet::Int32Type, int32_t)
+ TYPED_BATCH_CASE(INT64, Int64Type, ::parquet::Int64Type, int64_t)
+ TYPED_BATCH_CASE(FLOAT, FloatType, ::parquet::FloatType, float)
+ TYPED_BATCH_CASE(DOUBLE, DoubleType, ::parquet::DoubleType, double)
+ default:
+ return Status::NotImplemented(field_->type->ToString());
+ }
+}
+
+FlatColumnReader::FlatColumnReader(std::unique_ptr<Impl> impl) : impl_(std::move(impl)) {}
+
+FlatColumnReader::~FlatColumnReader() {}
+
+Status FlatColumnReader::NextBatch(int batch_size, std::shared_ptr<Array>* out) {
+ return impl_->NextBatch(batch_size, out);
+}
+
+} // namespace parquet
+} // namespace arrow