Skip to content

Commit 41248c8

Browse files
HuaHuaYHuaHuaY
committed
Implement basic parquet file rewriter
1 parent f32f8ac commit 41248c8

15 files changed

Lines changed: 1226 additions & 9 deletions

‎cpp/src/parquet/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,7 @@ set(PARQUET_SRCS
172172
encryption/internal_file_encryptor.cc
173173
exception.cc
174174
file_reader.cc
175+
file_rewriter.cc
175176
file_writer.cc
176177
geospatial/statistics.cc
177178
geospatial/util_internal.cc
@@ -406,6 +407,8 @@ add_parquet_test(arrow-reader-writer-test
406407

407408
add_parquet_test(arrow-index-test SOURCES arrow/index_test.cc)
408409

410+
add_parquet_test(arrow-rewriter-test SOURCES arrow/arrow_rewriter_test.cc)
411+
409412
add_parquet_test(arrow-internals-test SOURCES arrow/path_internal_test.cc
410413
arrow/reconstruct_internal_test.cc)
411414

‎cpp/src/parquet/arrow/arrow_rewriter_test.cc‎

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#include <gmock/gmock.h>
19+
#include <gtest/gtest.h>
20+
21+
#include <memory>
22+
23+
#include "arrow/io/memory.h"
24+
#include "arrow/testing/gtest_util.h"
25+
#include "parquet/arrow/reader.h"
26+
#include "parquet/arrow/test_util.h"
27+
#include "parquet/file_reader.h"
28+
#include "parquet/file_rewriter.h"
29+
#include "parquet/platform.h"
30+
#include "parquet/properties.h"
31+
#include "parquet/test_util.h"
32+
33+
using arrow::Table;
34+
using arrow::io::BufferReader;
35+
36+
namespace parquet::arrow {
37+
38+
TEST(ParquetRewriterTest, SimpleRoundTrip) {
39+
auto rewriter_properties =
40+
RewriterProperties::Builder()
41+
.writer_properties(
42+
WriterProperties::Builder().disable_write_page_index()->build())
43+
->build();
44+
45+
auto schema = ::arrow::schema(
46+
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
47+
48+
std::shared_ptr<Buffer> buffer;
49+
50+
WriteFile(rewriter_properties->writer_properties(),
51+
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"]])"}), buffer);
52+
53+
auto sink = CreateOutputStream();
54+
auto rewriter =
55+
ParquetFileRewriter::Open({{std::make_shared<BufferReader>(buffer)}}, sink,
56+
{{NULLPTR}}, NULLPTR, rewriter_properties);
57+
rewriter->Rewrite();
58+
rewriter->Close();
59+
60+
ASSERT_OK_AND_ASSIGN(auto out_buffer, sink->Finish());
61+
auto file_reader = ParquetFileReader::Open(std::make_shared<BufferReader>(out_buffer));
62+
ASSERT_OK_AND_ASSIGN(auto reader, FileReader::Make(::arrow::default_memory_pool(),
63+
std::move(file_reader)));
64+
65+
ASSERT_OK_AND_ASSIGN(auto table, reader->ReadTable());
66+
ASSERT_OK(table->ValidateFull());
67+
68+
auto expected_table = ::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"]])"});
69+
AssertTablesEqual(*expected_table, *table);
70+
}
71+
72+
TEST(ParquetRewriterTest, ConcatRoundTrip) {
73+
auto rewriter_properties =
74+
RewriterProperties::Builder()
75+
.writer_properties(
76+
WriterProperties::Builder().enable_write_page_index()->build())
77+
->build();
78+
79+
auto schema = ::arrow::schema(
80+
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
81+
82+
std::shared_ptr<Buffer> buffer_up;
83+
std::shared_ptr<Buffer> buffer_down;
84+
85+
WriteFile(rewriter_properties->writer_properties(),
86+
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"]])"}), buffer_up);
87+
WriteFile(rewriter_properties->writer_properties(),
88+
::arrow::TableFromJSON(schema, {R"([[3, "c"]])"}), buffer_down);
89+
90+
auto sink = CreateOutputStream();
91+
auto rewriter =
92+
ParquetFileRewriter::Open({{std::make_shared<BufferReader>(buffer_up),
93+
std::make_shared<BufferReader>(buffer_down)}},
94+
sink, {{NULLPTR, NULLPTR}}, NULLPTR, rewriter_properties);
95+
rewriter->Rewrite();
96+
rewriter->Close();
97+
98+
ASSERT_OK_AND_ASSIGN(auto out_buffer, sink->Finish());
99+
auto file_reader = ParquetFileReader::Open(std::make_shared<BufferReader>(out_buffer));
100+
ASSERT_OK_AND_ASSIGN(auto reader, FileReader::Make(::arrow::default_memory_pool(),
101+
std::move(file_reader)));
102+
103+
ASSERT_OK_AND_ASSIGN(auto table, reader->ReadTable());
104+
ASSERT_OK(table->ValidateFull());
105+
106+
auto expected_table =
107+
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"});
108+
AssertTablesEqual(*expected_table, *table);
109+
}
110+
111+
TEST(ParquetRewriterTest, JoinRoundTrip) {
112+
auto rewriter_properties =
113+
RewriterProperties::Builder()
114+
.writer_properties(
115+
WriterProperties::Builder().enable_write_page_index()->build())
116+
->build();
117+
118+
auto left_schema = ::arrow::schema(
119+
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
120+
auto right_schema = ::arrow::schema({::arrow::field("c", ::arrow::int64())});
121+
122+
std::shared_ptr<Buffer> buffer_left;
123+
std::shared_ptr<Buffer> buffer_right;
124+
125+
WriteFile(rewriter_properties->writer_properties(),
126+
::arrow::TableFromJSON(left_schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"}),
127+
buffer_left);
128+
WriteFile(rewriter_properties->writer_properties(),
129+
::arrow::TableFromJSON(right_schema, {R"([[10], [20], [30]])"}),
130+
buffer_right);
131+
132+
auto sink = CreateOutputStream();
133+
auto rewriter = ParquetFileRewriter::Open(
134+
{{std::make_shared<BufferReader>(buffer_left)},
135+
{std::make_shared<BufferReader>(buffer_right)}},
136+
sink, {{NULLPTR}, {NULLPTR}}, NULLPTR, rewriter_properties);
137+
rewriter->Rewrite();
138+
rewriter->Close();
139+
140+
ASSERT_OK_AND_ASSIGN(auto out_buffer, sink->Finish());
141+
auto file_reader = ParquetFileReader::Open(std::make_shared<BufferReader>(out_buffer));
142+
ASSERT_OK_AND_ASSIGN(auto reader, FileReader::Make(::arrow::default_memory_pool(),
143+
std::move(file_reader)));
144+
145+
ASSERT_OK_AND_ASSIGN(auto table, reader->ReadTable());
146+
ASSERT_OK(table->ValidateFull());
147+
148+
auto expected_schema = ::arrow::schema({::arrow::field("a", ::arrow::int32()),
149+
::arrow::field("b", ::arrow::utf8()),
150+
::arrow::field("c", ::arrow::int64())});
151+
auto expected_table = ::arrow::TableFromJSON(
152+
expected_schema, {R"([[1, "a", 10], [2, "b", 20], [3, "c", 30]])"});
153+
AssertTablesEqual(*expected_table, *table);
154+
}
155+
156+
TEST(ParquetRewriterTest, ConcatJoinRoundTrip) {
157+
auto rewriter_properties = RewriterProperties::Builder()
158+
.writer_properties(WriterProperties::Builder()
159+
.enable_write_page_index()
160+
->max_row_group_length(2)
161+
->build())
162+
->build();
163+
164+
auto left_schema = ::arrow::schema(
165+
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
166+
auto right_schema = ::arrow::schema({::arrow::field("c", ::arrow::int64())});
167+
168+
std::shared_ptr<Buffer> buffer_left_up;
169+
std::shared_ptr<Buffer> buffer_left_down;
170+
std::shared_ptr<Buffer> buffer_right;
171+
172+
WriteFile(rewriter_properties->writer_properties(),
173+
::arrow::TableFromJSON(left_schema, {R"([[1, "a"], [2, "b"]])"}),
174+
buffer_left_up);
175+
WriteFile(rewriter_properties->writer_properties(),
176+
::arrow::TableFromJSON(left_schema, {R"([[3, "c"]])"}), buffer_left_down);
177+
WriteFile(rewriter_properties->writer_properties(),
178+
::arrow::TableFromJSON(right_schema, {R"([[10], [20], [30]])"}),
179+
buffer_right);
180+
181+
auto sink = CreateOutputStream();
182+
auto rewriter = ParquetFileRewriter::Open(
183+
{{std::make_shared<BufferReader>(buffer_left_up),
184+
std::make_shared<BufferReader>(buffer_left_down)},
185+
{std::make_shared<BufferReader>(buffer_right)}},
186+
sink, {{NULLPTR, NULLPTR}, {NULLPTR}}, NULLPTR, rewriter_properties);
187+
rewriter->Rewrite();
188+
rewriter->Close();
189+
190+
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Buffer> out_buffer, sink->Finish());
191+
auto file_reader = ParquetFileReader::Open(std::make_shared<BufferReader>(out_buffer));
192+
ASSERT_OK_AND_ASSIGN(auto reader, FileReader::Make(::arrow::default_memory_pool(),
193+
std::move(file_reader)));
194+
195+
ASSERT_OK_AND_ASSIGN(std::shared_ptr<Table> table, reader->ReadTable());
196+
ASSERT_OK(table->ValidateFull());
197+
198+
auto expected_schema = ::arrow::schema({::arrow::field("a", ::arrow::int32()),
199+
::arrow::field("b", ::arrow::utf8()),
200+
::arrow::field("c", ::arrow::int64())});
201+
auto expected_table = ::arrow::TableFromJSON(
202+
expected_schema, {R"([[1, "a", 10], [2, "b", 20], [3, "c", 30]])"});
203+
AssertTablesEqual(*expected_table, *table);
204+
}
205+
206+
TEST(ParquetRewriterTest, JoinRowCountsMismatch) {
207+
auto rewriter_properties =
208+
RewriterProperties::Builder()
209+
.writer_properties(
210+
WriterProperties::Builder().enable_write_page_index()->build())
211+
->build();
212+
213+
auto schema1 = ::arrow::schema({::arrow::field("a", ::arrow::int32())});
214+
auto schema2 = ::arrow::schema({::arrow::field("b", ::arrow::int32())});
215+
216+
std::shared_ptr<Buffer> buffer1;
217+
std::shared_ptr<Buffer> buffer2;
218+
219+
WriteFile(rewriter_properties->writer_properties(),
220+
::arrow::TableFromJSON(schema1, {R"([[1], [2]])"}), buffer1);
221+
WriteFile(rewriter_properties->writer_properties(),
222+
::arrow::TableFromJSON(schema2, {R"([[3], [4], [5]])"}), buffer2);
223+
224+
auto sink = CreateOutputStream();
225+
226+
EXPECT_THROW_THAT(
227+
[&]() {
228+
ParquetFileRewriter::Open({{std::make_shared<BufferReader>(buffer1)},
229+
{std::make_shared<BufferReader>(buffer2)}},
230+
sink, {{NULLPTR}, {NULLPTR}}, NULLPTR,
231+
rewriter_properties);
232+
},
233+
ParquetException,
234+
::testing::Property(
235+
&ParquetException::what,
236+
::testing::HasSubstr("The number of rows in each block must match")));
237+
}
238+
239+
TEST(ParquetRewriterTest, InvalidInputDimensions) {
240+
auto rewriter_properties =
241+
RewriterProperties::Builder()
242+
.writer_properties(
243+
WriterProperties::Builder().enable_write_page_index()->build())
244+
->build();
245+
246+
auto schema = ::arrow::schema({::arrow::field("a", ::arrow::int32())});
247+
std::shared_ptr<Buffer> buffer;
248+
WriteFile(rewriter_properties->writer_properties(),
249+
::arrow::TableFromJSON(schema, {R"([[1]])"}), buffer);
250+
251+
auto sink = CreateOutputStream();
252+
253+
EXPECT_THROW_THAT(
254+
[&]() {
255+
ParquetFileRewriter::Open({{std::make_shared<BufferReader>(buffer)}}, sink, {},
256+
NULLPTR, rewriter_properties);
257+
},
258+
ParquetException,
259+
::testing::Property(
260+
&ParquetException::what,
261+
::testing::HasSubstr(
262+
"The number of sources and sources_metadata must be the same")));
263+
264+
EXPECT_THROW_THAT(
265+
[&]() {
266+
ParquetFileRewriter::Open({{std::make_shared<BufferReader>(buffer)}}, sink, {{}},
267+
NULLPTR, rewriter_properties);
268+
},
269+
ParquetException,
270+
::testing::Property(
271+
&ParquetException::what,
272+
::testing::HasSubstr(
273+
"The number of sources and sources_metadata must be the same")));
274+
}
275+
276+
} // namespace parquet::arrow

‎cpp/src/parquet/arrow/test_util.h‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,18 +28,23 @@
2828
#include "arrow/array/builder_binary.h"
2929
#include "arrow/array/builder_decimal.h"
3030
#include "arrow/array/builder_primitive.h"
31+
#include "arrow/table.h"
3132
#include "arrow/testing/gtest_util.h"
3233
#include "arrow/testing/random.h"
3334
#include "arrow/type_fwd.h"
3435
#include "arrow/type_traits.h"
3536
#include "arrow/util/decimal.h"
3637
#include "arrow/util/float16.h"
38+
#include "parquet/arrow/schema.h"
39+
#include "parquet/arrow/writer.h"
3740
#include "parquet/column_reader.h"
41+
#include "parquet/file_writer.h"
3842
#include "parquet/test_util.h"
3943

4044
namespace parquet {
4145

4246
using internal::RecordReader;
47+
using schema::GroupNode;
4348

4449
namespace arrow {
4550

@@ -482,6 +487,29 @@ void ExpectArrayT<::arrow::BooleanType>(void* expected, Array* result) {
482487
EXPECT_TRUE(result->Equals(*expected_array));
483488
}
484489

490+
void WriteFile(const std::shared_ptr<WriterProperties>& writer_properties,
491+
const std::shared_ptr<::arrow::Table>& table,
492+
std::shared_ptr<Buffer>& buffer) {
493+
// Get schema from table.
494+
auto schema = table->schema();
495+
std::shared_ptr<SchemaDescriptor> parquet_schema;
496+
auto arrow_writer_properties = default_arrow_writer_properties();
497+
ASSERT_OK_NO_THROW(ToParquetSchema(schema.get(), *writer_properties,
498+
*arrow_writer_properties, &parquet_schema));
499+
auto schema_node = std::static_pointer_cast<GroupNode>(parquet_schema->schema_root());
500+
501+
// Write table to buffer.
502+
auto sink = CreateOutputStream();
503+
auto pool = ::arrow::default_memory_pool();
504+
auto writer = ParquetFileWriter::Open(sink, schema_node, writer_properties);
505+
std::unique_ptr<FileWriter> arrow_writer;
506+
ASSERT_OK(FileWriter::Make(pool, std::move(writer), schema, arrow_writer_properties,
507+
&arrow_writer));
508+
ASSERT_OK_NO_THROW(arrow_writer->WriteTable(*table));
509+
ASSERT_OK_NO_THROW(arrow_writer->Close());
510+
ASSERT_OK_AND_ASSIGN(buffer, sink->Finish());
511+
}
512+
485513
} // namespace arrow
486514

487515
} // namespace parquet

‎cpp/src/parquet/bloom_filter_writer.cc‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,9 @@ class BloomFilterBuilderImpl : public BloomFilterBuilder {
162162

163163
BloomFilter* CreateBloomFilter(int32_t column_ordinal) override;
164164

165+
void InsertBloomFilter(int32_t column_ordinal,
166+
std::unique_ptr<BloomFilter> bloom_filter) override;
167+
165168
IndexLocations WriteTo(::arrow::io::OutputStream* sink) override;
166169

167170
private:
@@ -219,6 +222,26 @@ BloomFilter* BloomFilterBuilderImpl::CreateBloomFilter(int32_t column_ordinal) {
219222
return curr_rg_bfs.emplace(column_ordinal, std::move(bf)).first->second.get();
220223
}
221224

225+
void BloomFilterBuilderImpl::InsertBloomFilter(
226+
int32_t column_ordinal, std::unique_ptr<BloomFilter> bloom_filter) {
227+
auto opts = properties_->bloom_filter_options(schema_->Column(column_ordinal)->path());
228+
if (!opts.has_value() || bloom_filter == nullptr) {
229+
return;
230+
}
231+
232+
CheckState(column_ordinal);
233+
234+
auto& curr_rg_bfs = *bloom_filters_.rbegin();
235+
if (curr_rg_bfs.find(column_ordinal) != curr_rg_bfs.cend()) {
236+
std::stringstream ss;
237+
ss << "Bloom filter already exists for column: " << column_ordinal
238+
<< ", row group: " << (bloom_filters_.size() - 1);
239+
throw ParquetException(ss.str());
240+
}
241+
242+
curr_rg_bfs.emplace(column_ordinal, std::move(bloom_filter));
243+
}
244+
222245
IndexLocations BloomFilterBuilderImpl::WriteTo(::arrow::io::OutputStream* sink) {
223246
if (finished_) {
224247
throw ParquetException("Cannot write a finished BloomFilterBuilder");

0 commit comments

Comments
 (0)