Skip to content

Commit 57df653

Browse files
HuaHuaYHuaHuaY
committed
address review
1 parent a07aab0 commit 57df653

6 files changed

Lines changed: 82 additions & 81 deletions

File tree

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

Lines changed: 53 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,9 @@ TEST(ParquetRewriterTest, SimpleRoundTrip) {
4747
auto schema = ::arrow::schema(
4848
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
4949

50-
std::shared_ptr<Buffer> buffer;
51-
52-
WriteFile(rewriter_properties->writer_properties(),
53-
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"]])"}), buffer);
50+
ASSERT_OK_AND_ASSIGN(auto buffer, WriteFile(rewriter_properties->writer_properties(),
51+
::arrow::TableFromJSON(
52+
schema, {R"([[1, "a"], [2, "b"]])"})));
5453

5554
auto sink = CreateOutputStream();
5655
auto rewriter =
@@ -81,13 +80,13 @@ TEST(ParquetRewriterTest, ConcatRoundTrip) {
8180
auto schema = ::arrow::schema(
8281
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
8382

84-
std::shared_ptr<Buffer> buffer_up;
85-
std::shared_ptr<Buffer> buffer_down;
86-
87-
WriteFile(rewriter_properties->writer_properties(),
88-
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"]])"}), buffer_up);
89-
WriteFile(rewriter_properties->writer_properties(),
90-
::arrow::TableFromJSON(schema, {R"([[3, "c"]])"}), buffer_down);
83+
ASSERT_OK_AND_ASSIGN(
84+
auto buffer_up,
85+
WriteFile(rewriter_properties->writer_properties(),
86+
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"]])"})));
87+
ASSERT_OK_AND_ASSIGN(auto buffer_down,
88+
WriteFile(rewriter_properties->writer_properties(),
89+
::arrow::TableFromJSON(schema, {R"([[3, "c"]])"})));
9190

9291
auto sink = CreateOutputStream();
9392
auto rewriter =
@@ -121,15 +120,15 @@ TEST(ParquetRewriterTest, JoinRoundTrip) {
121120
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
122121
auto right_schema = ::arrow::schema({::arrow::field("c", ::arrow::int64())});
123122

124-
std::shared_ptr<Buffer> buffer_left;
125-
std::shared_ptr<Buffer> buffer_right;
126-
127-
WriteFile(rewriter_properties->writer_properties(),
128-
::arrow::TableFromJSON(left_schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"}),
129-
buffer_left);
130-
WriteFile(rewriter_properties->writer_properties(),
131-
::arrow::TableFromJSON(right_schema, {R"([[10], [20], [30]])"}),
132-
buffer_right);
123+
ASSERT_OK_AND_ASSIGN(
124+
auto buffer_left,
125+
WriteFile(
126+
rewriter_properties->writer_properties(),
127+
::arrow::TableFromJSON(left_schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"})));
128+
ASSERT_OK_AND_ASSIGN(
129+
auto buffer_right,
130+
WriteFile(rewriter_properties->writer_properties(),
131+
::arrow::TableFromJSON(right_schema, {R"([[10], [20], [30]])"})));
133132

134133
auto sink = CreateOutputStream();
135134
auto rewriter = ParquetFileRewriter::Open(
@@ -167,18 +166,17 @@ TEST(ParquetRewriterTest, ConcatJoinRoundTrip) {
167166
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
168167
auto right_schema = ::arrow::schema({::arrow::field("c", ::arrow::int64())});
169168

170-
std::shared_ptr<Buffer> buffer_left_up;
171-
std::shared_ptr<Buffer> buffer_left_down;
172-
std::shared_ptr<Buffer> buffer_right;
173-
174-
WriteFile(rewriter_properties->writer_properties(),
175-
::arrow::TableFromJSON(left_schema, {R"([[1, "a"], [2, "b"]])"}),
176-
buffer_left_up);
177-
WriteFile(rewriter_properties->writer_properties(),
178-
::arrow::TableFromJSON(left_schema, {R"([[3, "c"]])"}), buffer_left_down);
179-
WriteFile(rewriter_properties->writer_properties(),
180-
::arrow::TableFromJSON(right_schema, {R"([[10], [20], [30]])"}),
181-
buffer_right);
169+
ASSERT_OK_AND_ASSIGN(
170+
auto buffer_left_up,
171+
WriteFile(rewriter_properties->writer_properties(),
172+
::arrow::TableFromJSON(left_schema, {R"([[1, "a"], [2, "b"]])"})));
173+
ASSERT_OK_AND_ASSIGN(auto buffer_left_down,
174+
WriteFile(rewriter_properties->writer_properties(),
175+
::arrow::TableFromJSON(left_schema, {R"([[3, "c"]])"})));
176+
ASSERT_OK_AND_ASSIGN(
177+
auto buffer_right,
178+
WriteFile(rewriter_properties->writer_properties(),
179+
::arrow::TableFromJSON(right_schema, {R"([[10], [20], [30]])"})));
182180

183181
auto sink = CreateOutputStream();
184182
auto rewriter = ParquetFileRewriter::Open(
@@ -212,23 +210,23 @@ TEST(ParquetRewriterTest, JoinRowCountsMismatch) {
212210
WriterProperties::Builder().enable_write_page_index()->build())
213211
->build();
214212

215-
auto schema1 = ::arrow::schema({::arrow::field("a", ::arrow::int32())});
216-
auto schema2 = ::arrow::schema({::arrow::field("b", ::arrow::int32())});
217-
218-
std::shared_ptr<Buffer> buffer1;
219-
std::shared_ptr<Buffer> buffer2;
213+
auto left_schema = ::arrow::schema({::arrow::field("a", ::arrow::int32())});
214+
auto right_schema = ::arrow::schema({::arrow::field("b", ::arrow::int32())});
220215

221-
WriteFile(rewriter_properties->writer_properties(),
222-
::arrow::TableFromJSON(schema1, {R"([[1], [2]])"}), buffer1);
223-
WriteFile(rewriter_properties->writer_properties(),
224-
::arrow::TableFromJSON(schema2, {R"([[3], [4], [5]])"}), buffer2);
216+
ASSERT_OK_AND_ASSIGN(auto buffer_left,
217+
WriteFile(rewriter_properties->writer_properties(),
218+
::arrow::TableFromJSON(left_schema, {R"([[1], [2]])"})));
219+
ASSERT_OK_AND_ASSIGN(
220+
auto buffer_right,
221+
WriteFile(rewriter_properties->writer_properties(),
222+
::arrow::TableFromJSON(right_schema, {R"([[3], [4], [5]])"})));
225223

226224
auto sink = CreateOutputStream();
227225

228226
EXPECT_THROW_THAT(
229227
[&]() {
230-
ParquetFileRewriter::Open({{std::make_shared<BufferReader>(buffer1)},
231-
{std::make_shared<BufferReader>(buffer2)}},
228+
ParquetFileRewriter::Open({{std::make_shared<BufferReader>(buffer_left)},
229+
{std::make_shared<BufferReader>(buffer_right)}},
232230
sink, {{NULLPTR}, {NULLPTR}}, NULLPTR,
233231
rewriter_properties);
234232
},
@@ -246,9 +244,10 @@ TEST(ParquetRewriterTest, InvalidInputDimensions) {
246244
->build();
247245

248246
auto schema = ::arrow::schema({::arrow::field("a", ::arrow::int32())});
249-
std::shared_ptr<Buffer> buffer;
250-
WriteFile(rewriter_properties->writer_properties(),
251-
::arrow::TableFromJSON(schema, {R"([[1]])"}), buffer);
247+
248+
ASSERT_OK_AND_ASSIGN(auto buffer,
249+
WriteFile(rewriter_properties->writer_properties(),
250+
::arrow::TableFromJSON(schema, {R"([[1]])"})));
252251

253252
auto sink = CreateOutputStream();
254253

@@ -285,10 +284,10 @@ TEST(ParquetRewriterTest, AddCompression) {
285284
auto schema = ::arrow::schema(
286285
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
287286

288-
std::shared_ptr<Buffer> buffer;
289-
WriteFile(original_writer_props,
290-
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"}),
291-
buffer);
287+
ASSERT_OK_AND_ASSIGN(
288+
auto buffer,
289+
WriteFile(original_writer_props,
290+
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"})));
292291

293292
{
294293
auto file_reader = ParquetFileReader::Open(std::make_shared<BufferReader>(buffer));
@@ -344,10 +343,10 @@ TEST(ParquetRewriterTest, ChangeCompression) {
344343
auto schema = ::arrow::schema(
345344
{::arrow::field("a", ::arrow::int32()), ::arrow::field("b", ::arrow::utf8())});
346345

347-
std::shared_ptr<Buffer> buffer;
348-
WriteFile(original_writer_props,
349-
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"}),
350-
buffer);
346+
ASSERT_OK_AND_ASSIGN(
347+
auto buffer,
348+
WriteFile(original_writer_props,
349+
::arrow::TableFromJSON(schema, {R"([[1, "a"], [2, "b"], [3, "c"]])"})));
351350

352351
{
353352
auto file_reader = ParquetFileReader::Open(std::make_shared<BufferReader>(buffer));

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

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
#include "arrow/array/builder_binary.h"
2929
#include "arrow/array/builder_decimal.h"
3030
#include "arrow/array/builder_primitive.h"
31+
#include "arrow/result.h"
3132
#include "arrow/table.h"
3233
#include "arrow/testing/gtest_util.h"
3334
#include "arrow/testing/random.h"
@@ -487,27 +488,27 @@ void ExpectArrayT<::arrow::BooleanType>(void* expected, Array* result) {
487488
EXPECT_TRUE(result->Equals(*expected_array));
488489
}
489490

490-
void WriteFile(const std::shared_ptr<WriterProperties>& writer_properties,
491-
const std::shared_ptr<::arrow::Table>& table,
492-
std::shared_ptr<Buffer>& buffer) {
491+
::arrow::Result<std::shared_ptr<Buffer>> WriteFile(
492+
const std::shared_ptr<WriterProperties>& writer_properties,
493+
const std::shared_ptr<::arrow::Table>& table) {
493494
// Get schema from table.
494495
auto schema = table->schema();
495496
std::shared_ptr<SchemaDescriptor> parquet_schema;
496497
auto arrow_writer_properties = default_arrow_writer_properties();
497-
ASSERT_OK_NO_THROW(ToParquetSchema(schema.get(), *writer_properties,
498-
*arrow_writer_properties, &parquet_schema));
498+
RETURN_NOT_OK(ToParquetSchema(schema.get(), *writer_properties,
499+
*arrow_writer_properties, &parquet_schema));
499500
auto schema_node = std::static_pointer_cast<GroupNode>(parquet_schema->schema_root());
500501

501502
// Write table to buffer.
502503
auto sink = CreateOutputStream();
503504
auto pool = ::arrow::default_memory_pool();
504505
auto writer = ParquetFileWriter::Open(sink, schema_node, writer_properties);
505506
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());
507+
RETURN_NOT_OK(FileWriter::Make(pool, std::move(writer), schema, arrow_writer_properties,
508+
&arrow_writer));
509+
RETURN_NOT_OK(arrow_writer->WriteTable(*table));
510+
RETURN_NOT_OK(arrow_writer->Close());
511+
return sink->Finish();
511512
}
512513

513514
} // namespace arrow

‎cpp/src/parquet/file_rewriter.cc‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,11 @@
1818
#include "parquet/file_rewriter.h"
1919

2020
#include <algorithm>
21+
#include <format>
2122
#include <memory>
2223
#include <numeric>
2324
#include <optional>
2425
#include <ranges>
25-
#include <sstream>
2626
#include <unordered_set>
2727
#include <utility>
2828

@@ -465,10 +465,9 @@ class SingleFileRewriter {
465465
BloomFilterBuilder* bloom_filter_builder,
466466
int64_t& total_bytes_written) {
467467
if (current_row_group_index_ >= metadata_->num_row_groups()) {
468-
std::stringstream ss;
469-
ss << "Trying to read row group " << current_row_group_index_
470-
<< " but file only has " << metadata_->num_row_groups() << " row groups";
471-
throw ParquetException(ss.str());
468+
throw ParquetException("Trying to read row group ", current_row_group_index_,
469+
" but file only has ", metadata_->num_row_groups(),
470+
" row group");
472471
}
473472
auto row_group_reader = parquet_file_reader_->RowGroup(current_row_group_index_);
474473
auto page_index_reader = page_index_reader_
@@ -576,18 +575,18 @@ class JoinRewriter {
576575
for (size_t i = 1; i < rewriters_.size(); ++i) {
577576
if (auto current_row_counts = rewriters_[i]->row_group_row_counts();
578577
row_counts != current_row_counts) {
578+
// TODO(anyone): use `std::format("{}", row_counts)` instead when C++23 available
579579
auto vecToString = [](const std::vector<int64_t>& v) {
580580
if (v.empty()) {
581581
return std::string("[]");
582582
}
583-
std::ostringstream oss;
584-
oss << "[" << v[0];
585-
// TODO(anyone): use std::format and std::views::join_with when C++23 available
583+
584+
std::string s = std::format("[{}", v[0]);
586585
for (const auto& val : v | std::views::drop(1)) {
587-
oss << ", " << val;
586+
s += std::format(", {}", val);
588587
}
589-
oss << "]";
590-
return oss.str();
588+
s += "]";
589+
return s;
591590
};
592591
throw ParquetException(
593592
"The number of rows in each block must match! No.0 blocks row counts: ",
@@ -606,7 +605,7 @@ class JoinRewriter {
606605
auto path = schema_desc.Column(i)->path()->ToDotString();
607606
if (auto [_, inserted] = column_paths.emplace(path); !inserted) {
608607
// TODO(HuaHuaY): support choose one column from columns with same path
609-
throw ParquetException("NotImplemented, files have same column path: ", path);
608+
ParquetException::NYI(std::format("files have the same column path: {}", path));
610609
}
611610
}
612611

@@ -741,8 +740,7 @@ class GeneratedFile : public ParquetFileRewriter::Contents {
741740
// Unencrypted parquet files always start with PAR1
742741
PARQUET_THROW_NOT_OK(sink_->Write(kParquetMagic, 4));
743742
} else {
744-
throw ParquetException(
745-
"NotImplemented, rewriter does not support to write encrypted files.");
743+
ParquetException::NYI("Rewriter does not support writing encrypted files.");
746744
}
747745

748746
auto new_schema = rewriter_->schema().schema_root();

‎cpp/src/parquet/file_rewriter.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@ class PARQUET_EXPORT ParquetFileRewriter {
4747
void Open(std::unique_ptr<Contents> contents);
4848
void Close();
4949

50+
/// Rewrite all Parquet files.
51+
///
52+
/// This method may throw.
5053
void Rewrite();
5154

5255
private:

‎cpp/src/parquet/metadata.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2052,7 +2052,7 @@ ColumnChunkMetaDataBuilder* RowGroupMetaDataBuilder::NextColumnChunk() {
20522052

20532053
void RowGroupMetaDataBuilder::NextColumnChunk(
20542054
std::unique_ptr<ColumnChunkMetaData> cc_metadata, int64_t shift) {
2055-
return impl_->NextColumnChunk(std::move(cc_metadata), shift);
2055+
impl_->NextColumnChunk(std::move(cc_metadata), shift);
20562056
}
20572057

20582058
int RowGroupMetaDataBuilder::current_column() const { return impl_->current_column(); }

‎cpp/src/parquet/page_index.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -825,11 +825,11 @@ class PageIndexBuilderImpl final : public PageIndexBuilder {
825825

826826
WriteResult result;
827827

828-
/// Serialize column index ordered by row group ordinal and then column ordinal.
828+
// Serialize column index ordered by row group ordinal and then column ordinal.
829829
result.column_index_locations =
830830
SerializeIndex(column_index_builders_, column_indices_, sink);
831831

832-
/// Serialize offset index ordered by row group ordinal and then column ordinal.
832+
// Serialize offset index ordered by row group ordinal and then column ordinal.
833833
result.offset_index_locations =
834834
SerializeIndex(offset_index_builders_, offset_indices_, sink);
835835

0 commit comments

Comments
 (0)