Skip to content

Commit 7be60a0

Browse files
HuaHuaYHuaHuaY
committed
simplify RowGroupRewriter
1 parent faa4c8f commit 7be60a0

1 file changed

Lines changed: 90 additions & 49 deletions

File tree

‎cpp/src/parquet/file_rewriter.cc‎

Lines changed: 90 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -78,11 +78,81 @@ const std::shared_ptr<RewriterProperties>& default_rewriter_properties() {
7878
return default_rewriter_properties;
7979
}
8080

81+
class ColumnChunkRewriter {
82+
public:
83+
ColumnChunkRewriter(std::shared_ptr<ArrowInputFile> source,
84+
std::shared_ptr<ArrowOutputStream> sink,
85+
const RewriterProperties& props,
86+
std::unique_ptr<ColumnChunkMetaData> metadata,
87+
std::shared_ptr<RowGroupPageIndexReader> page_index_reader,
88+
std::shared_ptr<RowGroupBloomFilterReader> bloom_filter_reader,
89+
int column_ordinal)
90+
: source_(std::move(source)),
91+
sink_(std::move(sink)),
92+
props_(props),
93+
metadata_(std::move(metadata)),
94+
page_index_reader_(std::move(page_index_reader)),
95+
bloom_filter_reader_(std::move(bloom_filter_reader)),
96+
column_ordinal_(column_ordinal) {}
97+
98+
void WriteColumnChunkData(RowGroupMetaDataBuilder& rg_metadata_builder,
99+
PageIndexBuilder* page_index_builder,
100+
BloomFilterBuilder* bloom_filter_builder,
101+
int64_t& total_bytes_written) {
102+
// TODO(HuaHuaY): add else branch to rewrite column chunk with new encoding,
103+
// compression, etc.
104+
bool fast_copy = true;
105+
if (fast_copy) {
106+
PARQUET_ASSIGN_OR_THROW(int64_t sink_offset, sink_->Tell());
107+
int64_t shift = sink_offset - metadata_->start_offset();
108+
int64_t total_uncompressed_size = metadata_->total_uncompressed_size();
109+
int64_t total_compressed_size = metadata_->total_compressed_size();
110+
111+
auto stream = props_.reader_properties().GetStream(
112+
source_, metadata_->start_offset(), total_compressed_size);
113+
CopyStream(stream, sink_, total_compressed_size, props_.memory_pool());
114+
PARQUET_THROW_NOT_OK(stream->Close());
115+
116+
rg_metadata_builder.NextColumnChunk(std::move(metadata_), shift);
117+
118+
if (page_index_reader_ != nullptr && page_index_builder != nullptr) {
119+
auto column_index = page_index_reader_->GetColumnIndex(column_ordinal_);
120+
auto offset_index = page_index_reader_->GetOffsetIndex(column_ordinal_);
121+
if (column_index != nullptr) {
122+
page_index_builder->SetColumnIndex(column_ordinal_, column_index);
123+
}
124+
if (offset_index != nullptr) {
125+
page_index_builder->SetOffsetIndex(column_ordinal_, offset_index, shift);
126+
}
127+
}
128+
129+
if (bloom_filter_reader_ != nullptr && bloom_filter_builder != nullptr) {
130+
auto bloom_filter = bloom_filter_reader_->GetColumnBloomFilter(column_ordinal_);
131+
if (bloom_filter != nullptr) {
132+
bloom_filter_builder->InsertBloomFilter(column_ordinal_,
133+
std::move(bloom_filter));
134+
}
135+
}
136+
137+
total_bytes_written += total_uncompressed_size;
138+
}
139+
}
140+
141+
private:
142+
std::shared_ptr<ArrowInputFile> source_;
143+
std::shared_ptr<ArrowOutputStream> sink_;
144+
const RewriterProperties& props_;
145+
std::unique_ptr<ColumnChunkMetaData> metadata_;
146+
std::shared_ptr<RowGroupPageIndexReader> page_index_reader_;
147+
std::shared_ptr<RowGroupBloomFilterReader> bloom_filter_reader_;
148+
int column_ordinal_;
149+
};
150+
81151
class RowGroupRewriter {
82152
public:
83153
RowGroupRewriter(std::shared_ptr<ArrowInputFile> source,
84154
std::shared_ptr<ArrowOutputStream> sink,
85-
const RewriterProperties* props,
155+
const RewriterProperties& props,
86156
std::shared_ptr<RowGroupReader> row_group_reader,
87157
std::shared_ptr<RowGroupPageIndexReader> page_index_reader,
88158
std::shared_ptr<RowGroupBloomFilterReader> bloom_filter_reader)
@@ -94,26 +164,25 @@ class RowGroupRewriter {
94164
bloom_filter_reader_(std::move(bloom_filter_reader)),
95165
metadata_(row_group_reader_->metadata()) {}
96166

97-
void WriteRowGroupData(RowGroupMetaDataBuilder* rg_metadata_builder,
167+
void WriteRowGroupData(RowGroupMetaDataBuilder& rg_metadata_builder,
98168
PageIndexBuilder* page_index_builder,
99169
BloomFilterBuilder* bloom_filter_builder,
100170
int64_t& total_bytes_written) {
101-
rg_metadata_builder->set_num_rows(metadata_->num_rows());
171+
rg_metadata_builder.set_num_rows(metadata_->num_rows());
102172

103173
bool fast_copy = metadata_->file_offset() != 0;
104174
if (fast_copy) {
105175
PARQUET_ASSIGN_OR_THROW(int64_t sink_offset, sink_->Tell());
106176
int64_t shift = sink_offset - metadata_->file_offset();
107177

108-
auto stream = props_->reader_properties().GetStream(
178+
auto stream = props_.reader_properties().GetStream(
109179
source_, metadata_->file_offset(), metadata_->total_compressed_size());
110-
CopyStream(stream, sink_, metadata_->total_compressed_size(),
111-
props_->memory_pool());
180+
CopyStream(stream, sink_, metadata_->total_compressed_size(), props_.memory_pool());
112181
PARQUET_THROW_NOT_OK(stream->Close());
113182

114183
for (int i = 0; i < metadata_->num_columns(); ++i) {
115184
auto cc_metadata = metadata_->ColumnChunk(i);
116-
rg_metadata_builder->NextColumnChunk(std::move(cc_metadata), shift);
185+
rg_metadata_builder.NextColumnChunk(std::move(cc_metadata), shift);
117186

118187
auto column_index =
119188
page_index_reader_ ? page_index_reader_->GetColumnIndex(i) : nullptr;
@@ -138,45 +207,18 @@ class RowGroupRewriter {
138207
} else {
139208
for (int i = 0; i < metadata_->num_columns(); ++i) {
140209
auto cc_metadata = metadata_->ColumnChunk(i);
141-
142-
PARQUET_ASSIGN_OR_THROW(int64_t sink_offset, sink_->Tell());
143-
int64_t shift = sink_offset - cc_metadata->start_offset();
144-
145-
// TODO(HuaHuaY): add else branch to rewrite column chunk with new encoding,
146-
// compression, etc.
147-
if (true) {
148-
auto column_index = page_index_reader_->GetColumnIndex(i);
149-
auto offset_index = page_index_reader_->GetOffsetIndex(i);
150-
auto bloom_filter = bloom_filter_reader_->GetColumnBloomFilter(i);
151-
152-
auto stream = props_->reader_properties().GetStream(
153-
source_, cc_metadata->start_offset(), cc_metadata->total_compressed_size());
154-
CopyStream(stream, sink_, cc_metadata->total_compressed_size(),
155-
props_->memory_pool());
156-
PARQUET_THROW_NOT_OK(stream->Close());
157-
158-
rg_metadata_builder->NextColumnChunk(std::move(cc_metadata), shift);
159-
160-
if (column_index != nullptr && page_index_builder != nullptr) {
161-
page_index_builder->SetColumnIndex(i, column_index);
162-
}
163-
if (offset_index != nullptr && page_index_builder != nullptr) {
164-
page_index_builder->SetOffsetIndex(i, offset_index, shift);
165-
}
166-
if (bloom_filter != nullptr && bloom_filter_builder != nullptr) {
167-
bloom_filter_builder->InsertBloomFilter(i, std::move(bloom_filter));
168-
}
169-
170-
total_bytes_written += cc_metadata->total_uncompressed_size();
171-
}
210+
ColumnChunkRewriter rewriter(source_, sink_, props_, std::move(cc_metadata),
211+
page_index_reader_, bloom_filter_reader_, i);
212+
rewriter.WriteColumnChunkData(rg_metadata_builder, page_index_builder,
213+
bloom_filter_builder, total_bytes_written);
172214
}
173215
}
174216
}
175217

176218
private:
177219
std::shared_ptr<ArrowInputFile> source_;
178220
std::shared_ptr<ArrowOutputStream> sink_;
179-
const RewriterProperties* props_;
221+
const RewriterProperties& props_;
180222
std::shared_ptr<RowGroupReader> row_group_reader_;
181223
std::shared_ptr<RowGroupPageIndexReader> page_index_reader_;
182224
std::shared_ptr<RowGroupBloomFilterReader> bloom_filter_reader_;
@@ -188,12 +230,12 @@ class SingleFileRewriter {
188230
SingleFileRewriter(std::shared_ptr<ArrowInputFile> source,
189231
std::shared_ptr<ArrowOutputStream> sink,
190232
std::shared_ptr<FileMetaData> source_metadata,
191-
const RewriterProperties* props)
233+
const RewriterProperties& props)
192234
: source_(source),
193235
sink_(std::move(sink)),
194236
props_(props),
195237
parquet_file_reader_(ParquetFileReader::Open(
196-
std::move(source), props->reader_properties(), std::move(source_metadata))),
238+
std::move(source), props_.reader_properties(), std::move(source_metadata))),
197239
page_index_reader_(parquet_file_reader_->GetPageIndexReader()),
198240
bloom_filter_reader_(parquet_file_reader_->GetBloomFilterReader()),
199241
metadata_(parquet_file_reader_->metadata()) {
@@ -205,7 +247,7 @@ class SingleFileRewriter {
205247
{/*column_index=*/true, /*offset_index=*/true});
206248
}
207249

208-
void WriteRowGroupData(RowGroupMetaDataBuilder* rg_metadata_builder,
250+
void WriteRowGroupData(RowGroupMetaDataBuilder& rg_metadata_builder,
209251
PageIndexBuilder* page_index_builder,
210252
BloomFilterBuilder* bloom_filter_builder,
211253
int64_t& total_bytes_written) {
@@ -248,7 +290,7 @@ class SingleFileRewriter {
248290
private:
249291
std::shared_ptr<ArrowInputFile> source_;
250292
std::shared_ptr<ArrowOutputStream> sink_;
251-
const RewriterProperties* props_;
293+
const RewriterProperties& props_;
252294
std::unique_ptr<ParquetFileReader> parquet_file_reader_;
253295
std::shared_ptr<PageIndexReader> page_index_reader_;
254296
BloomFilterReader& bloom_filter_reader_;
@@ -268,7 +310,7 @@ class ConcatRewriter {
268310
}
269311
}
270312

271-
void WriteRowGroupData(RowGroupMetaDataBuilder* rg_metadata_builder,
313+
void WriteRowGroupData(RowGroupMetaDataBuilder& rg_metadata_builder,
272314
PageIndexBuilder* page_index_builder,
273315
BloomFilterBuilder* bloom_filter_builder,
274316
int64_t& total_bytes_written) {
@@ -358,7 +400,7 @@ class JoinRewriter {
358400
schema_.Init(new_root);
359401
}
360402

361-
void WriteRowGroupData(RowGroupMetaDataBuilder* rg_metadata_builder,
403+
void WriteRowGroupData(RowGroupMetaDataBuilder& rg_metadata_builder,
362404
PageIndexBuilder* page_index_builder,
363405
BloomFilterBuilder* bloom_filter_builder,
364406
int64_t& total_bytes_written) {
@@ -422,7 +464,7 @@ class GeneratedFile : public ParquetFileRewriter::Contents {
422464

423465
void Rewrite() override {
424466
while (rewriter_->HasMoreRowGroup()) {
425-
auto* rg_metadata_builder = metadata_builder_->AppendRowGroup();
467+
auto& rg_metadata_builder = *metadata_builder_->AppendRowGroup();
426468
if (page_index_builder_) {
427469
page_index_builder_->AppendRowGroup();
428470
}
@@ -432,7 +474,7 @@ class GeneratedFile : public ParquetFileRewriter::Contents {
432474
int64_t total_bytes_written = 0;
433475
rewriter_->WriteRowGroupData(rg_metadata_builder, page_index_builder_.get(),
434476
bloom_filter_builder_.get(), total_bytes_written);
435-
rg_metadata_builder->Finish(total_bytes_written);
477+
rg_metadata_builder.Finish(total_bytes_written);
436478
}
437479
if (page_index_builder_) {
438480
page_index_builder_->Finish();
@@ -469,8 +511,7 @@ class GeneratedFile : public ParquetFileRewriter::Contents {
469511
concat_rewriters.reserve(sources[i].size());
470512
for (size_t j = 0; j < sources[i].size(); ++j) {
471513
concat_rewriters.emplace_back(std::make_unique<SingleFileRewriter>(
472-
std::move(sources[i][j]), sink_, std::move(sources_metadata[i][j]),
473-
props_.get()));
514+
std::move(sources[i][j]), sink_, std::move(sources_metadata[i][j]), *props_));
474515
}
475516
rewriters.emplace_back(
476517
std::make_unique<ConcatRewriter>(std::move(concat_rewriters)));

0 commit comments

Comments
 (0)