def get(self): """ Get batched sampled `Edges`. Return: An `Edges` object, shape=[`batch_size`] Raise: `graphlearn.OutOfRangeError` """ mask_type = utils.get_mask_type(self._edge_type, self._mask) state = self._graph.edge_state.get(mask_type) req = pywrap.new_get_edges_request(mask_type, self._strategy, self._batch_size, state) res = pywrap.new_get_edges_response() status = self._client.get_edges(req, res) if not status.ok(): self._graph.edge_state.inc(mask_type) else: src_ids = pywrap.get_edge_src_id(res) dst_ids = pywrap.get_edge_dst_id(res) edge_ids = pywrap.get_edge_id(res) pywrap.del_op_response(res) pywrap.del_op_request(req) errors.raise_exception_on_not_ok_status(status) edges = self._graph.get_edges(self._edge_type, src_ids, dst_ids) edges.edge_ids = edge_ids return edges
def node(self, source, node_type, decoder, option=None, mask=utils.Mask.NONE): """ Add graph vertices that will be loaded from a given path. Args: source (string): Data source path where to load the nodes. node_type (string): Indicates the type of the added nodes. decoder (Decoder): A Decoder object to describe the data schema. mask (TRAIN | TEST | VAL): Mark the source as TRAIN data, TEST data or VAL data for the given node_type in the graph. """ if not isinstance(source, str): raise ValueError('source for node() must be string.') if not isinstance(node_type, str): raise ValueError('node_type for node() must be string.') if not isinstance(decoder, data.Decoder): raise ValueError( 'decoder must be an instance of `Decoder`, got {}'.format( type(decoder))) node_type = utils.get_mask_type(node_type, mask) self._node_decoders[node_type] = decoder node_source = self._construct_node_source(source, node_type, decoder, option) self._node_sources.append(node_source) return self
def get(self): """ Get batched sampled nodes. Return: A `Nodes` object, shape=[`batch_size`] Raise: `graphlearn.OutOfRangeError` """ mask_type = utils.get_mask_type(self._type, self._mask) state = self._graph.node_state.get(mask_type) req = pywrap.new_get_nodes_request(mask_type, self._strategy, self._node_from, self._batch_size, state) res = pywrap.new_get_nodes_response() status = self._client.get_nodes(req, res) if not status.ok(): self._graph.node_state.inc(mask_type) else: ids = pywrap.get_node_ids(res) pywrap.del_op_response(res) pywrap.del_op_request(req) errors.raise_exception_on_not_ok_status(status) nodes = self._graph.get_nodes(self._node_type, ids) return nodes
def _lazy_E(self, t, mask=utils.Mask.NONE): dag = gsl.Dag(self) params = {pywrap.kEdgeType: utils.get_mask_type(t, mask)} source_node = gsl.TraverseSourceEdgeDagNode(dag, op_name="GetEdges", params=params) source_node.set_output_field(pywrap.kEdgeIds) source_node.set_path(t, pywrap.NodeFrom.NODE) # Add sink node to dag gsl.SinkNode(dag) return source_node
def _lazy_V(self, t, node_from=pywrap.NodeFrom.NODE, mask=utils.Mask.NONE): dag = gsl.Dag(self) params = { pywrap.kNodeType: utils.get_mask_type(t, mask), pywrap.kNodeFrom: int(node_from) } source_node = gsl.TraverseVertexDagNode(dag, op_name="GetNodes", params=params) source_node.set_output_field(pywrap.kNodeIds) source_node.set_path(t, node_from) # Add sink node to dag gsl.SinkNode(dag) return source_node
def V(self, t, feed=None, node_from=pywrap.NodeFrom.NODE, mask=utils.Mask.NONE): """ Entry of GSL, starting from VERTEX. Args: t (string): The type of node which is the entry of query or the type of edge when node is from edge source or dst. feed (None| numpy.ndarray | types.GeneratorType | `Nodes`): When `feed` is not `None`, the `type` should be a node type, which means query the attributes of the specified node ids. None: Default. Sample nodes with the following .shuffle and .batch API. numpy.ndarray: Any shape of ids. Get nodes of the given ids and node_type. types.Generator: A generator of numpy.ndarray. Get nodes of generated ids and given node_type. `Nodes`: A `Nodes` object. node_from (NodeFrom): Default is `NodeFrom.NODE`, which means sample or or iterate node from node. `NodeFrom.EDGE_SRC` means sample or iterate node from source node of edge, and `NodeFrom.EDGE_DST` means sample or iterate node from destination node of edge. If node is from edge, the `type` must be an edge type. mask (NONE | TRAIN | TEST | VAL): The given node set is indexed by both the raw node type and mask value. The default mask value is NONE, which plays nothing on the index. """ if feed is not None: raise NotImplementedError("`feed` is not supported for now.") dag = gsl.Dag(self) params = { pywrap.kNodeType: utils.get_mask_type(t, mask), pywrap.kNodeFrom: int(node_from) } source_node = gsl.TraverseVertexDagNode(dag, op_name="GetNodes", params=params) source_node.set_output_field(pywrap.kNodeIds) source_node.set_path(t, node_from) # Add sink node to dag gsl.SinkNode(dag) return source_node
def edge(self, source, edge_type, decoder=None, directed=True, option=None, mask=utils.Mask.NONE): """ Add graph edges that will be loaded from a given path. Args: source (string): Data source path where to load the edges. edge_type (tuple): A tuple of (src_type, dst_type, edge_type) that indicates types of the edges. decoder (Decoder): A Decoder object to describe the data schema. directed (boolean): Whether edges are directed. mask (TRAIN | TEST | VAL): Mark the source as TRAIN data, TEST data or VAL data for the given edge_type in the graph. """ if not isinstance(source, str): raise ValueError('source for edge() must be a string.') if not isinstance(edge_type, tuple) or len(edge_type) != 3: raise ValueError("edge_type for edge() must be a tuple of " "(src_type, dst_tye, edge_type).") if not decoder: decoder = data.Decoder() if not isinstance(decoder, data.Decoder): raise ValueError( 'decoder must be an instance of Decoder, got {}'.format( type(decoder))) masked_edge_type = utils.get_mask_type(edge_type[2], mask) self._edge_decoders[masked_edge_type] = decoder self._topology.add(masked_edge_type, edge_type[0], edge_type[1]) edge_source = self._construct_edge_source( source, (edge_type[0], edge_type[1], masked_edge_type), decoder, direction=pywrap.Direction.ORIGIN, option=option) self._edge_sources.append(edge_source) if not directed: self.add_reverse_edges(edge_type, source, decoder, option) return self