diff --git a/ml_metadata/proto/metadata_store.proto b/ml_metadata/proto/metadata_store.proto index 53dbb3e75..672b42d80 100644 --- a/ml_metadata/proto/metadata_store.proto +++ b/ml_metadata/proto/metadata_store.proto @@ -462,6 +462,9 @@ message Association { message ParentContext { optional int64 child_id = 1; optional int64 parent_id = 2; + + // Output only. + optional google.protobuf.Any system_metadata = 3; } // A self-contained provenance (sub)graph representation consists of MLMD nodes diff --git a/ml_metadata/util/record_parsing_utils.cc b/ml_metadata/util/record_parsing_utils.cc index 163c923ae..d2bd9f298 100644 --- a/ml_metadata/util/record_parsing_utils.cc +++ b/ml_metadata/util/record_parsing_utils.cc @@ -125,6 +125,14 @@ absl::Status ParseRecordSetToEdgeArray( CustomColumnParser()); } +absl::Status ParseRecordSetToEdgeArray( + const RecordSet& record_set, + std::vector& output_parent_contexts, + const CustomColumnParser& parser) { + return ParseRecordSetToMessageArray(record_set, output_parent_contexts, + parser); +} + absl::Status ParseNodeRecordSetToDedupedTypes( const RecordSet& node_record_set, std::vector& output_artifact_types, diff --git a/ml_metadata/util/record_parsing_utils.h b/ml_metadata/util/record_parsing_utils.h index ef0405049..0ad7dc49b 100644 --- a/ml_metadata/util/record_parsing_utils.h +++ b/ml_metadata/util/record_parsing_utils.h @@ -81,6 +81,14 @@ absl::Status ParseRecordSetToEdgeArray( absl::Status ParseRecordSetToEdgeArray( const RecordSet& record_set, std::vector& output_attributions); +// Converts `record_set` to a ParentContext array. +// Returns OK and the parsed result is outputted by `output_parent_contexts`. +// Returns error when internal error happens. +absl::Status ParseRecordSetToEdgeArray( + const RecordSet& record_set, + std::vector& output_parent_contexts, + const CustomColumnParser& parser = CustomColumnParser()); + // Extracts ArtifactType information from `node_record_set` to a deduped // ArtifactType array. // Returns OK and the parsed result is outputted by `output_artifact_types`. diff --git a/ml_metadata/util/record_parsing_utils_test.cc b/ml_metadata/util/record_parsing_utils_test.cc index b2f12e90a..f58ee58fd 100644 --- a/ml_metadata/util/record_parsing_utils_test.cc +++ b/ml_metadata/util/record_parsing_utils_test.cc @@ -29,6 +29,7 @@ namespace { using ::ml_metadata::testing::EqualsProto; using ::ml_metadata::testing::ParseTextProtoOrDie; using ::testing::ElementsAre; +using ::testing::IsEmpty; TEST(ParseRecordSetTest, ParseRecordSetToArtifactArraySuccess) { RecordSet record_set = ParseTextProtoOrDie( @@ -273,6 +274,40 @@ TEST(ParseRecordSetTest, ParseRecordSetToAttributionArraySuccess) { )pb")))); } +TEST(ParseRecordSetTest, ParseRecordSetToParentContextArraySuccess) { + RecordSet record_set = ParseTextProtoOrDie( + R"pb( + column_names: 'parent_id' + column_names: 'child_id' + records { values: '123' values: '321' } + records { values: '456' values: '654' } + )pb"); + + std::vector parent_contexts; + absl::Status status = ParseRecordSetToEdgeArray(record_set, parent_contexts); + EXPECT_EQ(status, absl::OkStatus()); + EXPECT_THAT(parent_contexts, + ElementsAre(EqualsProto(ParseTextProtoOrDie(R"pb( + parent_id: 123 child_id: 321 + )pb")), + EqualsProto(ParseTextProtoOrDie(R"pb( + parent_id: 456 child_id: 654 + )pb")))); +} + +TEST(ParseRecordSetTest, ParseRecordSetToParentContextArrayEmptySuccess) { + RecordSet record_set = ParseTextProtoOrDie( + R"pb( + column_names: 'parent_id' + column_names: 'child_id' + )pb"); + + std::vector parent_contexts; + absl::Status status = ParseRecordSetToEdgeArray(record_set, parent_contexts); + EXPECT_EQ(status, absl::OkStatus()); + EXPECT_THAT(parent_contexts, IsEmpty()); +} + TEST(ParseRecordSetTest, MismatchRecordAndFieldNameIgnored) { RecordSet record_set = ParseTextProtoOrDie( R"pb( @@ -457,6 +492,7 @@ TEST(ParseRecordSetTest, ParseRecordSetToDedupedContextTypeArraySuccess) { name: 'context_type_1' )pb")))); } + } // namespace } // namespace testing } // namespace ml_metadata