| 32 | namespace parquet::arrow { |
| 33 | |
| 34 | TEST(Metadata, AppendMetadata) { |
| 35 | // A sample table, type and structure does not matter in this test case |
| 36 | auto schema = ::arrow::schema({::arrow::field("f", ::arrow::utf8())}); |
| 37 | auto table = ::arrow::Table::Make( |
| 38 | schema, {::arrow::ArrayFromJSON(::arrow::utf8(), R"(["a", "b", "c"])")}); |
| 39 | |
| 40 | auto sink = CreateOutputStream(); |
| 41 | ArrowWriterProperties::Builder builder; |
| 42 | builder.store_schema(); |
| 43 | ASSERT_OK_AND_ASSIGN(auto writer, |
| 44 | parquet::arrow::FileWriter::Open( |
| 45 | *schema, ::arrow::default_memory_pool(), sink, |
| 46 | parquet::default_writer_properties(), builder.build())); |
| 47 | |
| 48 | auto kv_meta = std::make_shared<KeyValueMetadata>(); |
| 49 | kv_meta->Append("test_key_1", "test_value_1"); |
| 50 | // <test_key_2, test_value_2_temp> would be overwritten later. |
| 51 | kv_meta->Append("test_key_2", "test_value_2_temp"); |
| 52 | ASSERT_OK(writer->AddKeyValueMetadata(kv_meta)); |
| 53 | |
| 54 | // Key value metadata that will be added to the file. |
| 55 | auto kv_meta_added = std::make_shared<::arrow::KeyValueMetadata>(); |
| 56 | kv_meta_added->Append("test_key_2", "test_value_2"); |
| 57 | kv_meta_added->Append("test_key_3", "test_value_3"); |
| 58 | |
| 59 | ASSERT_OK(writer->AddKeyValueMetadata(kv_meta_added)); |
| 60 | ASSERT_OK(writer->Close()); |
| 61 | |
| 62 | // return error if the file is closed |
| 63 | ASSERT_RAISES(IOError, writer->AddKeyValueMetadata(kv_meta_added)); |
| 64 | |
| 65 | auto verify_key_value_metadata = |
| 66 | [&](const std::shared_ptr<const KeyValueMetadata>& key_value_metadata) { |
| 67 | ASSERT_TRUE(nullptr != key_value_metadata); |
| 68 | |
| 69 | // Verify keys that were added before file writer was closed are present. |
| 70 | for (int i = 1; i <= 3; ++i) { |
| 71 | auto index = std::to_string(i); |
| 72 | PARQUET_ASSIGN_OR_THROW(auto value, |
| 73 | key_value_metadata->Get("test_key_" + index)); |
| 74 | EXPECT_EQ("test_value_" + index, value); |
| 75 | } |
| 76 | EXPECT_TRUE(key_value_metadata->Contains("ARROW:schema")); |
| 77 | }; |
| 78 | // verify the metadata in writer |
| 79 | verify_key_value_metadata(writer->metadata()->key_value_metadata()); |
| 80 | |
| 81 | ASSERT_OK(writer->Close()); |
| 82 | |
| 83 | ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); |
| 84 | // verify the metadata in reader |
| 85 | { |
| 86 | std::unique_ptr<FileReader> reader; |
| 87 | FileReaderBuilder reader_builder; |
| 88 | ASSERT_OK_NO_THROW( |
| 89 | reader_builder.Open(std::make_shared<::arrow::io::BufferReader>(buffer))); |
| 90 | ASSERT_OK( |
| 91 | reader_builder.properties(default_arrow_reader_properties())->Build(&reader)); |
nothing calls this directly
no test coverage detected