"examples/vscode:/vscode.git/clone" did not exist on "269983dbcd221acb64d03817fb25f0e27788bbaf"
graph_serialize.h 2.91 KB
Newer Older
VoVAllen's avatar
VoVAllen committed
1
2
/*!
 *  Copyright (c) 2019 by Contributors
3
 * \file graph/serialize/graph_serialize.h
VoVAllen's avatar
VoVAllen committed
4
5
 * \brief Graph serialization header
 */
6
7
#ifndef DGL_GRAPH_SERIALIZE_GRAPH_SERIALIZE_H_
#define DGL_GRAPH_SERIALIZE_GRAPH_SERIALIZE_H_
VoVAllen's avatar
VoVAllen committed
8
9

#include <dgl/array.h>
10
#include <dgl/graph.h>
VoVAllen's avatar
VoVAllen committed
11
#include <dgl/immutable_graph.h>
12
#include <dgl/packed_func_ext.h>
VoVAllen's avatar
VoVAllen committed
13
#include <dgl/runtime/container.h>
14
#include <dgl/runtime/ndarray.h>
VoVAllen's avatar
VoVAllen committed
15
#include <dgl/runtime/object.h>
16
17
18
19
#include <dmlc/io.h>
#include <dmlc/type_traits.h>

#include <algorithm>
VoVAllen's avatar
VoVAllen committed
20
#include <iostream>
21
#include <memory>
VoVAllen's avatar
VoVAllen committed
22
23
#include <string>
#include <utility>
24
25
#include <vector>

26
#include "../../c_api_common.h"
27
28
#include "dglgraph_data.h"
#include "heterograph_data.h"
VoVAllen's avatar
VoVAllen committed
29
30

using dgl::ImmutableGraph;
31
using dgl::runtime::NDArray;
VoVAllen's avatar
VoVAllen committed
32
33
34
35
36
using namespace dgl::runtime;

namespace dgl {
namespace serialize {

37
38
39
40
enum GraphType : uint64_t {
  kMutableGraph = 0ull,
  kImmutableGraph = 1ull,
  kHeteroGraph = 2ull
VoVAllen's avatar
VoVAllen committed
41
42
};

43
constexpr uint64_t kDGLSerializeMagic = 0xDD2E4FF046B4A13F;
VoVAllen's avatar
VoVAllen committed
44
45
46

class StorageMetaDataObject : public runtime::Object {
 public:
47
  // For saving DGLGraph
VoVAllen's avatar
VoVAllen committed
48
49
50
51
52
  dgl_id_t num_graph;
  Value nodes_num_list;
  Value edges_num_list;
  Map<std::string, Value> labels_list;
  List<GraphData> graph_data;
53

VoVAllen's avatar
VoVAllen committed
54
55
  static constexpr const char *_type_key = "graph_serialize.StorageMetaData";

56
  void SetMetaData(dgl_id_t num_graph, std::vector<int64_t> nodes_num_list,
VoVAllen's avatar
VoVAllen committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
                   std::vector<int64_t> edges_num_list,
                   std::vector<NamedTensor> labels_list);

  void SetGraphData(std::vector<GraphData> gdata);

  void VisitAttrs(AttrVisitor *v) final {
    v->Visit("num_graph", &num_graph);
    v->Visit("nodes_num_list", &nodes_num_list);
    v->Visit("edges_num_list", &edges_num_list);
    v->Visit("labels", &labels_list);
    v->Visit("graph_data", &graph_data);
  }

  DGL_DECLARE_OBJECT_TYPE_INFO(StorageMetaDataObject, runtime::Object);
};

class StorageMetaData : public runtime::ObjectRef {
 public:
75
76
  DGL_DEFINE_OBJECT_REF_METHODS(StorageMetaData, runtime::ObjectRef,
                                StorageMetaDataObject);
VoVAllen's avatar
VoVAllen committed
77
78
79
80
81
82
83

  /*! \brief create a new StorageMetaData reference */
  static StorageMetaData Create() {
    return StorageMetaData(std::make_shared<StorageMetaDataObject>());
  }
};

84
85
86
87
88
89
StorageMetaData LoadDGLGraphFiles(const std::string &filename,
                                  std::vector<dgl_id_t> idx_list,
                                  bool onlyMeta);

StorageMetaData LoadDGLGraphs(const std::string &filename,
                              std::vector<dgl_id_t> idx_list, bool onlyMeta);
VoVAllen's avatar
VoVAllen committed
90

91
bool SaveDGLGraphs(std::string filename, List<GraphData> graph_data,
VoVAllen's avatar
VoVAllen committed
92
93
                   std::vector<NamedTensor> labels_list);

94
95
std::vector<HeteroGraphData> LoadHeteroGraphs(const std::string &filename,
                                              std::vector<dgl_id_t> idx_list);
VoVAllen's avatar
VoVAllen committed
96
97
98
99
100
101

ImmutableGraphPtr ToImmutableGraph(GraphPtr g);

}  // namespace serialize
}  // namespace dgl

102
#endif  // DGL_GRAPH_SERIALIZE_GRAPH_SERIALIZE_H_