diff --git a/cpp/src/graphar/high-level/edges_builder.h b/cpp/src/graphar/high-level/edges_builder.h index ce847f4e3..382d0b96a 100644 --- a/cpp/src/graphar/high-level/edges_builder.h +++ b/cpp/src/graphar/high-level/edges_builder.h @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -280,6 +281,30 @@ class EdgesBuilder { return Status::OK(); } + /** + * @brief Add a property to all edges in the collection. + * + * @param property name of the property + * @param values vector of values where values[i] is mapped to the i-th edge + * in chunk-major order with size equal to the edges collection + * + * @return Status: ok or Status::Invalid error. + */ + [[nodiscard]] Status AddPropertyColumn(const std::string& property, + const std::vector& values) { + if (static_cast(values.size()) != num_edges_) { + return Status::Invalid( + "The size of values vector is not equal to the number of edges."); + } + + IdType value = 0; + for (auto& [id, edges] : edges_) { + for (Edge& edge : edges) { + edge.AddProperty(property, values[value++]); + } + } + return Status::OK(); + } /** * @brief Get the current number of edges in the collection. * @@ -471,7 +496,7 @@ class EdgesBuilder { std::shared_ptr edge_info_; std::string prefix_; AdjListType adj_list_type_; - std::unordered_map> edges_; + std::map> edges_; IdType vertex_chunk_size_; IdType num_vertices_; IdType num_edges_; diff --git a/cpp/src/graphar/high-level/vertices_builder.h b/cpp/src/graphar/high-level/vertices_builder.h index 38632a223..019df6f05 100644 --- a/cpp/src/graphar/high-level/vertices_builder.h +++ b/cpp/src/graphar/high-level/vertices_builder.h @@ -346,6 +346,28 @@ class VerticesBuilder { return Status::OK(); } + /** + * @brief Add a property to all vertices in the collection. + * + * @param property name of the property + * @param values vector of values where values[i] is mapped to the i-th vertex + * in insertion order with size equal to the vertices collection + * + * @return Status: ok or Status::Invalid error. + */ + [[nodiscard]] Status AddPropertyColumn(const std::string& property, + const std::vector& values) { + if (static_cast(values.size()) != num_vertices_) { + return Status::Invalid( + "The size of values vector is not equal to the number of vertices."); + } + + for (size_t i = 0; i < vertices_.size(); i++) { + vertices_[i].AddProperty(property, values[i]); + } + return Status::OK(); + } + /** * @brief Get the current number of vertices in the collection. * diff --git a/cpp/test/test_builder.cc b/cpp/test/test_builder.cc index 9a25b5af1..67026e264 100644 --- a/cpp/test/test_builder.cc +++ b/cpp/test/test_builder.cc @@ -119,6 +119,23 @@ TEST_CASE_METHOD(GlobalFixture, "Test_vertices_builder") { // check the number of vertices in builder REQUIRE(builder->GetNum() == lines); + // add property column + std::vector int_values(builder->GetNum()); + std::vector string_values(builder->GetNum()); + for (IdType i = 0; i < builder->GetNum(); i++) { + int_values[i] = i + 10; + string_values[i] = std::to_string(i); + } + + REQUIRE(builder->AddPropertyColumn("id", int_values).ok()); + REQUIRE(builder->AddPropertyColumn("firstName", string_values).ok()); + + int_values.push_back(100); + string_values.push_back("test"); + + REQUIRE(builder->AddPropertyColumn("id", int_values).IsInvalid()); + REQUIRE(builder->AddPropertyColumn("firstName", string_values).IsInvalid()); + // dump to files REQUIRE(builder->Dump().ok()); @@ -145,6 +162,37 @@ TEST_CASE_METHOD(GlobalFixture, "Test_vertices_builder") { auto row_group_meta = parquet_metadata->RowGroup(0); auto col_meta = row_group_meta->ColumnChunk(0); REQUIRE(col_meta->compression() == parquet::Compression::LZ4); + + // check that properties were added correctly + auto name_file = "/tmp/vertex/person/firstName_lastName_gender/chunk0"; + + std::unique_ptr name_reader; + + REQUIRE(graphar::util::OpenParquetArrowReader( + name_file, arrow::default_memory_pool(), &name_reader) + .ok()); + + auto id_col = parquet_table->GetColumnByName("id"); + + auto maybe_name_table = name_reader->ReadTable(); + REQUIRE(maybe_name_table.ok()); + auto name_table = maybe_name_table.ValueOrDie(); + auto name_col = name_table->GetColumnByName("firstName"); + + REQUIRE(name_col != nullptr); + REQUIRE(id_col != nullptr); + + REQUIRE(id_col->type()->id() == arrow::Type::INT64); + auto id_array = std::static_pointer_cast(id_col->chunk(0)); + auto name_array = + std::static_pointer_cast(name_col->chunk(0)); + + for (IdType i = 0; i < id_array->length(); i++) { + REQUIRE(id_array->Value(i) == i + 10); + } + for (IdType i = 0; i < name_array->length(); i++) { + REQUIRE(name_array->GetString(i) == std::to_string(i)); + } } TEST_CASE_METHOD(GlobalFixture, "test_edges_builder") { @@ -221,6 +269,17 @@ TEST_CASE_METHOD(GlobalFixture, "test_edges_builder") { // check the number of edges in builder REQUIRE(builder->GetNum() == lines); + // add property column + std::vector string_values(builder->GetNum(), + std::string("test_edge")); + + REQUIRE(builder->AddPropertyColumn("creationDate", string_values).ok()); + + string_values.push_back(std::string("test")); + + REQUIRE( + builder->AddPropertyColumn("creationDate", string_values).IsInvalid()); + // dump to files REQUIRE(builder->Dump().ok()); @@ -251,5 +310,17 @@ TEST_CASE_METHOD(GlobalFixture, "test_edges_builder") { auto row_group_meta = parquet_metadata->RowGroup(0); auto col_meta = row_group_meta->ColumnChunk(0); REQUIRE(col_meta->compression() == parquet::Compression::LZ4); + // check that properties were added correctly + + auto date_col = parquet_table->GetColumnByName("creationDate"); + + REQUIRE(date_col != nullptr); + + auto string_array = + std::static_pointer_cast(date_col->chunk(0)); + + for (IdType i = 0; i < string_array->length(); i++) { + REQUIRE(string_array->GetString(i) == "test_edge"); + } } } // namespace graphar