Skip to content

Commit 6312724

Browse files
wesmxhochy
authored andcommitted
PARQUET-830: Add parquet::arrow::OpenFile with additional properties and metadata args
I also slightly refactored the test suite to use OpenFile rather than using the `ParquetFileReader` ctor directly (`OpenFile` wasn't being used in the test suite). Needed for ARROW-471 Author: Wes McKinney <wes.mckinney@twosigma.com> Closes apache#219 from wesm/PARQUET-830 and squashes the following commits: bd17192 [Wes McKinney] Add parquet::arrow::OpenFile with additional properties and metadata arguments Change-Id: Ib00d04a9284b2108377a9ffac9faf8514b9e46cf
1 parent 2cbd797 commit 6312724

3 files changed

Lines changed: 46 additions & 20 deletions

File tree

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

Lines changed: 28 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -203,16 +203,17 @@ class TestParquetIO : public ::testing::Test {
203203
return ParquetFileWriter::Open(sink_, schema);
204204
}
205205

206-
std::unique_ptr<ParquetFileReader> ReaderFromSink() {
206+
void ReaderFromSink(std::unique_ptr<FileReader>* out) {
207207
std::shared_ptr<Buffer> buffer = sink_->GetBuffer();
208-
return ParquetFileReader::Open(std::make_shared<BufferReader>(buffer));
208+
ASSERT_OK_NO_THROW(OpenFile(std::make_shared<BufferReader>(buffer),
209+
::arrow::default_memory_pool(), ::parquet::default_reader_properties(),
210+
nullptr, out));
209211
}
210212

211213
void ReadSingleColumnFile(
212-
std::unique_ptr<ParquetFileReader> file_reader, std::shared_ptr<Array>* out) {
213-
FileReader reader(::arrow::default_memory_pool(), std::move(file_reader));
214+
std::unique_ptr<FileReader> file_reader, std::shared_ptr<Array>* out) {
214215
std::unique_ptr<FlatColumnReader> column_reader;
215-
ASSERT_OK_NO_THROW(reader.GetFlatColumn(0, &column_reader));
216+
ASSERT_OK_NO_THROW(file_reader->GetFlatColumn(0, &column_reader));
216217
ASSERT_NE(nullptr, column_reader.get());
217218

218219
ASSERT_OK(column_reader->NextBatch(SMALL_SIZE, out));
@@ -221,20 +222,24 @@ class TestParquetIO : public ::testing::Test {
221222

222223
void ReadAndCheckSingleColumnFile(::arrow::Array* values) {
223224
std::shared_ptr<::arrow::Array> out;
224-
ReadSingleColumnFile(ReaderFromSink(), &out);
225+
226+
std::unique_ptr<FileReader> reader;
227+
ReaderFromSink(&reader);
228+
ReadSingleColumnFile(std::move(reader), &out);
225229
ASSERT_TRUE(values->Equals(out));
226230
}
227231

228232
void ReadTableFromFile(
229-
std::unique_ptr<ParquetFileReader> file_reader, std::shared_ptr<Table>* out) {
230-
FileReader reader(::arrow::default_memory_pool(), std::move(file_reader));
231-
ASSERT_OK_NO_THROW(reader.ReadFlatTable(out));
233+
std::unique_ptr<FileReader> reader, std::shared_ptr<Table>* out) {
234+
ASSERT_OK_NO_THROW(reader->ReadFlatTable(out));
232235
ASSERT_NE(nullptr, out->get());
233236
}
234237

235238
void ReadAndCheckSingleColumnTable(const std::shared_ptr<::arrow::Array>& values) {
236239
std::shared_ptr<::arrow::Table> out;
237-
ReadTableFromFile(ReaderFromSink(), &out);
240+
std::unique_ptr<FileReader> reader;
241+
ReaderFromSink(&reader);
242+
ReadTableFromFile(std::move(reader), &out);
238243
ASSERT_EQ(1, out->num_columns());
239244
ASSERT_EQ(values->length(), out->num_rows());
240245

@@ -287,7 +292,9 @@ TYPED_TEST(TestParquetIO, SingleColumnTableRequiredWrite) {
287292
this->sink_, values->length(), default_writer_properties()));
288293

289294
std::shared_ptr<Table> out;
290-
this->ReadTableFromFile(this->ReaderFromSink(), &out);
295+
std::unique_ptr<FileReader> reader;
296+
this->ReaderFromSink(&reader);
297+
this->ReadTableFromFile(std::move(reader), &out);
291298
ASSERT_EQ(1, out->num_columns());
292299
ASSERT_EQ(100, out->num_rows());
293300

@@ -368,7 +375,9 @@ TYPED_TEST(TestParquetIO, SingleColumnTableRequiredChunkedWriteArrowIO) {
368375

369376
auto source = std::make_shared<BufferReader>(pbuffer);
370377
std::shared_ptr<::arrow::Table> out;
371-
this->ReadTableFromFile(ParquetFileReader::Open(std::move(source)), &out);
378+
std::unique_ptr<FileReader> reader;
379+
ASSERT_OK_NO_THROW(OpenFile(source, ::arrow::default_memory_pool(), &reader));
380+
this->ReadTableFromFile(std::move(reader), &out);
372381
ASSERT_EQ(1, out->num_columns());
373382
ASSERT_EQ(values->length(), out->num_rows());
374383

@@ -530,7 +539,9 @@ TEST_F(TestStringParquetIO, EmptyStringColumnRequiredWrite) {
530539
this->sink_, values->length(), default_writer_properties()));
531540

532541
std::shared_ptr<Table> out;
533-
this->ReadTableFromFile(this->ReaderFromSink(), &out);
542+
std::unique_ptr<FileReader> reader;
543+
this->ReaderFromSink(&reader);
544+
this->ReadTableFromFile(std::move(reader), &out);
534545
ASSERT_EQ(1, out->num_columns());
535546
ASSERT_EQ(100, out->num_rows());
536547

@@ -558,7 +569,7 @@ class TestPrimitiveParquetIO : public TestParquetIO<TestType> {
558569
typedef typename c_type_trait<TestType>::ArrowCType T;
559570

560571
void MakeTestFile(std::vector<T>& values, int num_chunks,
561-
std::unique_ptr<ParquetFileReader>* file_reader) {
572+
std::unique_ptr<FileReader>* reader) {
562573
std::shared_ptr<GroupNode> schema = this->MakeSchema(Repetition::REQUIRED);
563574
std::unique_ptr<ParquetFileWriter> file_writer = this->MakeWriter(schema);
564575
size_t chunk_size = values.size() / num_chunks;
@@ -578,12 +589,12 @@ class TestPrimitiveParquetIO : public TestParquetIO<TestType> {
578589
row_group_writer->Close();
579590
}
580591
file_writer->Close();
581-
*file_reader = this->ReaderFromSink();
592+
this->ReaderFromSink(reader);
582593
}
583594

584595
void CheckSingleColumnRequiredTableRead(int num_chunks) {
585596
std::vector<T> values(SMALL_SIZE, test_traits<TestType>::value);
586-
std::unique_ptr<ParquetFileReader> file_reader;
597+
std::unique_ptr<FileReader> file_reader;
587598
ASSERT_NO_THROW(MakeTestFile(values, num_chunks, &file_reader));
588599

589600
std::shared_ptr<Table> out;
@@ -598,7 +609,7 @@ class TestPrimitiveParquetIO : public TestParquetIO<TestType> {
598609

599610
void CheckSingleColumnRequiredRead(int num_chunks) {
600611
std::vector<T> values(SMALL_SIZE, test_traits<TestType>::value);
601-
std::unique_ptr<ParquetFileReader> file_reader;
612+
std::unique_ptr<FileReader> file_reader;
602613
ASSERT_NO_THROW(MakeTestFile(values, num_chunks, &file_reader));
603614

604615
std::shared_ptr<Array> out;

cpp/src/parquet/arrow/reader.cc

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -190,14 +190,22 @@ FileReader::~FileReader() {}
190190

191191
// Static ctor
192192
Status OpenFile(const std::shared_ptr<::arrow::io::ReadableFileInterface>& file,
193-
MemoryPool* allocator, std::unique_ptr<FileReader>* reader) {
194-
// TODO(wesm): reader properties
193+
MemoryPool* allocator, const ReaderProperties& props,
194+
const std::shared_ptr<FileMetaData>& metadata, std::unique_ptr<FileReader>* reader) {
195+
std::unique_ptr<RandomAccessSource> io_wrapper(new ArrowInputFile(file));
195196
std::unique_ptr<ParquetReader> pq_reader;
196-
PARQUET_CATCH_NOT_OK(pq_reader = ParquetReader::Open(file));
197+
PARQUET_CATCH_NOT_OK(
198+
pq_reader = ParquetReader::Open(std::move(io_wrapper), props, metadata));
197199
reader->reset(new FileReader(allocator, std::move(pq_reader)));
198200
return Status::OK();
199201
}
200202

203+
Status OpenFile(const std::shared_ptr<::arrow::io::ReadableFileInterface>& file,
204+
MemoryPool* allocator, std::unique_ptr<FileReader>* reader) {
205+
return OpenFile(file, allocator, ::parquet::default_reader_properties(),
206+
nullptr, reader);
207+
}
208+
201209
Status FileReader::GetFlatColumn(int i, std::unique_ptr<FlatColumnReader>* out) {
202210
return impl_->GetFlatColumn(i, out);
203211
}

cpp/src/parquet/arrow/reader.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,13 @@ class PARQUET_EXPORT FlatColumnReader {
139139

140140
// Helper function to create a file reader from an implementation of an Arrow
141141
// readable file
142+
//
143+
// metadata : separately-computed file metadata, can be nullptr
144+
PARQUET_EXPORT
145+
::arrow::Status OpenFile(const std::shared_ptr<::arrow::io::ReadableFileInterface>& file,
146+
::arrow::MemoryPool* allocator, const ReaderProperties& properties,
147+
const std::shared_ptr<FileMetaData>& metadata, std::unique_ptr<FileReader>* reader);
148+
142149
PARQUET_EXPORT
143150
::arrow::Status OpenFile(const std::shared_ptr<::arrow::io::ReadableFileInterface>& file,
144151
::arrow::MemoryPool* allocator, std::unique_ptr<FileReader>* reader);

0 commit comments

Comments
 (0)