basic_fragment_loader.h

Go to the documentation of this file.

 
#ifndef STORAGES_RT_MUTABLE_GRAPH_LOADER_BASIC_FRAGMENT_LOADER_H_
#define STORAGES_RT_MUTABLE_GRAPH_LOADER_BASIC_FRAGMENT_LOADER_H_
 
#include "flex/storages/rt_mutable_graph/file_names.h"
#include "flex/storages/rt_mutable_graph/mutable_property_fragment.h"
#include "flex/storages/rt_mutable_graph/schema.h"
 
namespace gs {
 
template <typename EDATA_T>
TypedMutableCsrBase<EDATA_T>* create_typed_csr(EdgeStrategy es,
 PropertyType edge_property) {
 if (es == EdgeStrategy::kSingle) {
 return new SingleMutableCsr<EDATA_T>(edge_property);
 } else if (es == EdgeStrategy::kMultiple) {
 return new MutableCsr<EDATA_T>(edge_property);
 } else if (es == EdgeStrategy::kNone) {
 return new EmptyCsr<EDATA_T>();
 }
 LOG(FATAL) << "not support edge strategy or edge data type";
 return nullptr;
}
 
enum class LoadingStatus {
 kLoading = 0,
 kLoaded = 1,
 kCommited = 2,
 kUnknown = 3,
};
 
// define << and >> for LoadingStatus
std::ostream& operator<<(std::ostream& os, const LoadingStatus& status);
 
std::istream& operator>>(std::istream& is, LoadingStatus& status);
 
// FragmentLoader should use this BasicFragmentLoader to construct
// mutable_csr_fragment.
class BasicFragmentLoader {
 public:
 BasicFragmentLoader(const Schema& schema, const std::string& prefix);
 
 void LoadFragment();
 
 // props vector is column_num X batch_size
 void AddVertexBatch(label_t v_label, const std::vector<vid_t>& vids,
 const std::vector<std::vector<Any>>& props);
 
 inline void SetVertexProperty(label_t v_label, size_t col_ind, vid_t vid,
 Any&& prop) {
 auto& table = vertex_data_[v_label];
 auto dst_columns = table.column_ptrs();
 CHECK(col_ind < dst_columns.size());
 dst_columns[col_ind]->set_any(vid, prop);
 }
#ifndef USE_PTHASH
 template <typename KEY_T>
 void FinishAddingVertex(label_t v_label,
 const IdIndexer<KEY_T, vid_t>& indexer) {
 CHECK(v_label < vertex_label_num_);
 std::string filename =
 vertex_map_prefix(schema_.get_vertex_label_name(v_label));
 auto primary_keys = schema_.get_vertex_primary_key(v_label);
 auto type = std::get<0>(primary_keys[0]);
 
 build_lf_indexer<KEY_T, vid_t>(
 indexer, LFIndexer<vid_t>::prefix() + "_" + filename,
 lf_indexers_[v_label], snapshot_dir(work_dir_, 0), tmp_dir(work_dir_),
 type);
 append_vertex_loading_progress(schema_.get_vertex_label_name(v_label),
 LoadingStatus::kLoaded);
 auto& v_data = vertex_data_[v_label];
 auto label_name = schema_.get_vertex_label_name(v_label);
 v_data.resize(lf_indexers_[v_label].size());
 v_data.dump(vertex_table_prefix(label_name), snapshot_dir(work_dir_, 0));
 append_vertex_loading_progress(label_name, LoadingStatus::kCommited);
 }
#else
 template <typename KEY_T>
 void FinishAddingVertex(label_t v_label,
 PTIndexerBuilder<KEY_T, vid_t>& indexer_builder) {
 CHECK(v_label < vertex_label_num_);
 std::string filename =
 vertex_map_prefix(schema_.get_vertex_label_name(v_label));
 indexer_builder.finish(PTIndexer<vid_t>::prefix() + "_" + filename,
 snapshot_dir(work_dir_, 0), lf_indexers_[v_label]);
 append_vertex_loading_progress(schema_.get_vertex_label_name(v_label),
 LoadingStatus::kLoaded);
 auto& v_data = vertex_data_[v_label];
 auto label_name = schema_.get_vertex_label_name(v_label);
 v_data.resize(lf_indexers_[v_label].size());
 v_data.dump(vertex_table_prefix(label_name), snapshot_dir(work_dir_, 0));
 append_vertex_loading_progress(label_name, LoadingStatus::kCommited);
 }
#endif
 
 template <typename EDATA_T>
 void AddNoPropEdgeBatch(label_t src_label_id, label_t dst_label_id,
 label_t edge_label_id) {
 size_t index = src_label_id * vertex_label_num_ * edge_label_num_ +
 dst_label_id * edge_label_num_ + edge_label_id;
 CHECK(ie_[index] == NULL);
 CHECK(oe_[index] == NULL);
 auto src_label_name = schema_.get_vertex_label_name(src_label_id);
 auto dst_label_name = schema_.get_vertex_label_name(dst_label_id);
 auto edge_label_name = schema_.get_edge_label_name(edge_label_id);
 EdgeStrategy oe_strategy = schema_.get_outgoing_edge_strategy(
 src_label_name, dst_label_name, edge_label_name);
 EdgeStrategy ie_strategy = schema_.get_incoming_edge_strategy(
 src_label_name, dst_label_name, edge_label_name);
 if constexpr (std::is_same_v<EDATA_T, std::string_view>) {
 const auto& prop = schema_.get_edge_properties(src_label_id, dst_label_id,
 edge_label_id);
 
 size_t max_length = PropertyType::STRING_DEFAULT_MAX_LENGTH;
 if (prop[0].IsVarchar()) {
 max_length = prop[0].additional_type_info.max_length;
 }
 dual_csr_list_[index] =
 new DualCsr<std::string_view>(oe_strategy, ie_strategy, max_length);
 } else {
 bool oe_mutable = schema_.outgoing_edge_mutable(
 src_label_name, dst_label_name, edge_label_name);
 bool ie_mutable = schema_.incoming_edge_mutable(
 src_label_name, dst_label_name, edge_label_name);
 dual_csr_list_[index] = new DualCsr<EDATA_T>(oe_strategy, ie_strategy,
 oe_mutable, ie_mutable);
 }
 ie_[index] = dual_csr_list_[index]->GetInCsr();
 oe_[index] = dual_csr_list_[index]->GetOutCsr();
 dual_csr_list_[index]->BatchInit(
 oe_prefix(src_label_name, dst_label_name, edge_label_name),
 ie_prefix(src_label_name, dst_label_name, edge_label_name),
 edata_prefix(src_label_name, dst_label_name, edge_label_name),
 tmp_dir(work_dir_), {}, {});
 }
 
 template <typename EDATA_T>
 static decltype(auto) get_casted_dual_csr(DualCsrBase* dual_csr) {
 if constexpr (std::is_same_v<EDATA_T, RecordView>) {
 auto casted_dual_csr = dynamic_cast<DualCsr<RecordView>*>(dual_csr);
 CHECK(casted_dual_csr != NULL);
 return casted_dual_csr;
 } else {
 auto casted_dual_csr = dynamic_cast<DualCsr<EDATA_T>*>(dual_csr);
 CHECK(casted_dual_csr != NULL);
 return casted_dual_csr;
 }
 }
 template <typename EDATA_T, typename VECTOR_T>
 void PutEdges(label_t src_label_id, label_t dst_label_id,
 label_t edge_label_id, const std::vector<VECTOR_T>& edges_vec,
 const std::vector<int32_t>& ie_degree,
 const std::vector<int32_t>& oe_degree, bool build_csr_in_mem) {
 size_t index = src_label_id * vertex_label_num_ * edge_label_num_ +
 dst_label_id * edge_label_num_ + edge_label_id;
 auto dual_csr = dual_csr_list_[index];
 CHECK(dual_csr != NULL);
 auto casted_dual_csr = get_casted_dual_csr<EDATA_T>(dual_csr);
 CHECK(casted_dual_csr != NULL);
 auto& src_indexer = lf_indexers_[src_label_id];
 auto& dst_indexer = lf_indexers_[dst_label_id];
 auto src_label_name = schema_.get_vertex_label_name(src_label_id);
 auto dst_label_name = schema_.get_vertex_label_name(dst_label_id);
 auto edge_label_name = schema_.get_edge_label_name(edge_label_id);
 
 auto INVALID_VID = std::numeric_limits<vid_t>::max();
 std::atomic<size_t> edge_count(0);
 if constexpr (std::is_same_v<EDATA_T, std::string_view>) {
 CHECK(ie_degree.size() == dst_indexer.size());
 CHECK(oe_degree.size() == src_indexer.size());
 if (build_csr_in_mem) {
 dual_csr->BatchInitInMemory(
 edata_prefix(src_label_name, dst_label_name, edge_label_name),
 tmp_dir(work_dir_), oe_degree, ie_degree);
 } else {
 dual_csr->BatchInit(
 oe_prefix(src_label_name, dst_label_name, edge_label_name),
 ie_prefix(src_label_name, dst_label_name, edge_label_name),
 edata_prefix(src_label_name, dst_label_name, edge_label_name),
 tmp_dir(work_dir_), oe_degree, ie_degree);
 }
 std::vector<std::thread> work_threads;
 for (size_t i = 0; i < edges_vec.size(); ++i) {
 work_threads.emplace_back(
 [&](int idx) {
 edge_count.fetch_add(edges_vec[idx].size());
 for (auto& edge : edges_vec[idx]) {
 if (std::get<1>(edge) == INVALID_VID ||
 std::get<0>(edge) == INVALID_VID) {
 VLOG(10) << "Skip invalid edge:" << std::get<0>(edge) << "->"
 << std::get<1>(edge);
 continue;
 }
 casted_dual_csr->BatchPutEdge(
 std::get<0>(edge), std::get<1>(edge), std::get<2>(edge));
 }
 },
 i);
 }
 for (auto& t : work_threads) {
 t.join();
 }
 
 append_edge_loading_progress(src_label_name, dst_label_name,
 edge_label_name, LoadingStatus::kLoaded);
 if (schema_.get_sort_on_compaction(src_label_name, dst_label_name,
 edge_label_name)) {
 dual_csr->SortByEdgeData(1);
 }
 dual_csr->Dump(
 oe_prefix(src_label_name, dst_label_name, edge_label_name),
 ie_prefix(src_label_name, dst_label_name, edge_label_name),
 edata_prefix(src_label_name, dst_label_name, edge_label_name),
 snapshot_dir(work_dir_, 0));
 } else {
 CHECK(ie_degree.size() == dst_indexer.size());
 CHECK(oe_degree.size() == src_indexer.size());
 
 if (build_csr_in_mem) {
 dual_csr->BatchInitInMemory(
 edata_prefix(src_label_name, dst_label_name, edge_label_name),
 tmp_dir(work_dir_), oe_degree, ie_degree);
 } else {
 dual_csr->BatchInit(
 oe_prefix(src_label_name, dst_label_name, edge_label_name),
 ie_prefix(src_label_name, dst_label_name, edge_label_name),
 edata_prefix(src_label_name, dst_label_name, edge_label_name),
 tmp_dir(work_dir_), oe_degree, ie_degree);
 }
 
 std::vector<std::thread> work_threads;
 for (size_t i = 0; i < edges_vec.size(); ++i) {
 work_threads.emplace_back(
 [&](int idx) {
 edge_count.fetch_add(edges_vec[idx].size());
 for (auto& edge : edges_vec[idx]) {
 if (std::get<1>(edge) == INVALID_VID ||
 std::get<0>(edge) == INVALID_VID) {
 VLOG(10) << "Skip invalid edge:" << std::get<0>(edge) << "->"
 << std::get<1>(edge);
 continue;
 }
 casted_dual_csr->BatchPutEdge(
 std::get<0>(edge), std::get<1>(edge), std::get<2>(edge));
 }
 },
 i);
 }
 for (auto& t : work_threads) {
 t.join();
 }
 append_edge_loading_progress(src_label_name, dst_label_name,
 edge_label_name, LoadingStatus::kLoaded);
 if (schema_.get_sort_on_compaction(src_label_name, dst_label_name,
 edge_label_name)) {
 dual_csr->SortByEdgeData(1);
 }
 
 dual_csr->Dump(
 oe_prefix(src_label_name, dst_label_name, edge_label_name),
 ie_prefix(src_label_name, dst_label_name, edge_label_name),
 edata_prefix(src_label_name, dst_label_name, edge_label_name),
 snapshot_dir(work_dir_, 0));
 }
 append_edge_loading_progress(src_label_name, dst_label_name,
 edge_label_name, LoadingStatus::kCommited);
 VLOG(10) << "Finish adding edge batch of size: " << edge_count.load();
 }
 
 Table& GetVertexTable(size_t ind) {
 CHECK(ind < vertex_data_.size());
 return vertex_data_[ind];
 }
 
 // get lf_indexer
 const IndexerType& GetLFIndexer(label_t v_label) const;
 IndexerType& GetLFIndexer(label_t v_label);
 
 const std::string& work_dir() const { return work_dir_; }
 
 void set_csr(label_t src_label_id, label_t dst_label_id,
 label_t edge_label_id, DualCsrBase* dual_csr);
 
 DualCsrBase* get_csr(label_t src_label_id, label_t dst_label_id,
 label_t edge_label_id);
 
 void init_edge_table(label_t src_label_id, label_t dst_label_id,
 label_t edge_label_id);
 
 private:
 // create status files for each vertex label and edge triplet pair.
 void append_vertex_loading_progress(const std::string& label_name,
 LoadingStatus status);
 
 void append_edge_loading_progress(const std::string& src_label_name,
 const std::string& dst_label_name,
 const std::string& edge_label_name,
 LoadingStatus status);
 void init_loading_status_file();
 void init_vertex_data();
 const Schema& schema_;
 std::string work_dir_;
 size_t vertex_label_num_, edge_label_num_;
 std::vector<IndexerType> lf_indexers_;
 std::vector<CsrBase*> ie_, oe_;
 std::vector<DualCsrBase*> dual_csr_list_;
 std::vector<Table> vertex_data_;
 
 // loading progress related
 std::mutex loading_progress_mutex_;
};
} // namespace gs
 
#endif // STORAGES_RT_MUTABLE_GRAPH_LOADER_BASIC_FRAGMENT_LOADER_H_