apply suggestions

mapleFU · May 28, 2024 · a9155c8 · a9155c8
1 parent 93b3bbb
commit a9155c8
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 11 deletions.
diff --git a/cpp/src/parquet/arrow/arrow_metadata_test.cc b/cpp/src/parquet/arrow/arrow_metadata_test.cc
@@ -23,6 +23,7 @@
 
 #include "parquet/api/writer.h"
 
+#include "parquet/arrow/reader.h"
 #include "parquet/arrow/schema.h"
 #include "parquet/arrow/writer.h"
 #include "parquet/file_writer.h"
@@ -46,7 +47,8 @@ TEST(Metadata, AppendMetadata) {
 
   auto kv_meta = std::make_shared<KeyValueMetadata>();
   kv_meta->Append("test_key_1", "test_value_1");
-  kv_meta->Append("test_key_2", "test_value_2_");
+  // <test_key_2, test_value_2_temp> would be overwritten later.
+  kv_meta->Append("test_key_2", "test_value_2_temp");
   ASSERT_OK(writer->AddKeyValueMetadata(kv_meta));
 
   // Key value metadata that will be added to the file.
@@ -60,16 +62,36 @@ TEST(Metadata, AppendMetadata) {
   // return error if the file is closed
   ASSERT_RAISES(IOError, writer->AddKeyValueMetadata(kv_meta_added));
 
-  const auto& key_value_metadata = writer->metadata()->key_value_metadata();
-  ASSERT_TRUE(nullptr != key_value_metadata);
+  auto verify_key_value_metadata =
+      [&](const std::shared_ptr<const KeyValueMetadata>& key_value_metadata) {
+        ASSERT_TRUE(nullptr != key_value_metadata);
 
-  // Verify keys that were added before file writer was closed are present.
-  for (int i = 1; i <= 3; ++i) {
-    auto index = std::to_string(i);
-    PARQUET_ASSIGN_OR_THROW(auto value, key_value_metadata->Get("test_key_" + index));
-    EXPECT_EQ("test_value_" + index, value);
+        // Verify keys that were added before file writer was closed are present.
+        for (int i = 1; i <= 3; ++i) {
+          auto index = std::to_string(i);
+          PARQUET_ASSIGN_OR_THROW(auto value,
+                                  key_value_metadata->Get("test_key_" + index));
+          EXPECT_EQ("test_value_" + index, value);
+        }
+        EXPECT_TRUE(key_value_metadata->Contains("ARROW:schema"));
+      };
+  // verify the metadata in writer
+  verify_key_value_metadata(writer->metadata()->key_value_metadata());
+
+  ASSERT_OK(writer->Close());
+
+  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());
+  // verify the metadata in reader
+  {
+    std::unique_ptr<FileReader> reader;
+    FileReaderBuilder reader_builder;
+    ASSERT_OK_NO_THROW(
+        reader_builder.Open(std::make_shared<::arrow::io::BufferReader>(buffer)));
+    ASSERT_OK(
+        reader_builder.properties(default_arrow_reader_properties())->Build(&reader));
+
+    verify_key_value_metadata(reader->parquet_reader()->metadata()->key_value_metadata());
   }
-  EXPECT_TRUE(key_value_metadata->Contains("ARROW:schema"));
 }
 
 }  // namespace parquet::arrow
diff --git a/cpp/src/parquet/arrow/writer.h b/cpp/src/parquet/arrow/writer.h
@@ -150,7 +150,7 @@ class PARQUET_EXPORT FileWriter {
   ///
   /// WARNING: If `store_schema` is enabled, `ARROW:schema` would be stored
   /// in the key-value metadata. Overwriting this key would result in
-  /// `store_schema` unusable during read.
+  /// `store_schema` being unusable during read.
   virtual ::arrow::Status AddKeyValueMetadata(
       const std::shared_ptr<const ::arrow::KeyValueMetadata>& key_value_metadata) = 0;
   /// \brief Return the file metadata, only available after calling Close().

diff --git a/cpp/src/parquet/file_writer.h b/cpp/src/parquet/file_writer.h
@@ -202,7 +202,7 @@ class PARQUET_EXPORT ParquetFileWriter {
 
   /// \brief Add key-value metadata to the file.
   /// \param[in] key_value_metadata the metadata to add.
-  /// \note This will overwrite any existing metadata with the same key.
+  /// \note This will overwrite any existing metadata with the same key(s).
   /// \throw ParquetException if Close() has been called.
   void AddKeyValueMetadata(
       const std::shared_ptr<const KeyValueMetadata>& key_value_metadata);