def get(self):
""" Get batched sampled nodes.
Return:
A `Nodes` object, shape=[`batch_size`]
Raise:
`graphlearn.OutOfRangeError`
"""
state = self._graph.node_state.get(self._type)
req = pywrap.new_get_node_req(self._type, self._strategy,
self._node_from, self._batch_size, state)
res = pywrap.new_get_node_res()
status = self._client.get_nodes(req, res)
if not status.ok():
self._graph.edge_state.inc(self._type)
else:
ids = pywrap.get_node_node_id_res(res)
pywrap.del_get_node_res(res)
pywrap.del_get_node_req(req)
raise_exception_on_not_ok_status(status)
nodes = self._graph.get_nodes(self._node_type, ids)
return nodes
def lookup_nodes(self, node_type, ids):
""" Get all the node properties.
"""
req = pywrap.new_lookup_nodes_req(node_type)
pywrap.set_lookup_nodes_req(req, ids)
res = pywrap.new_lookup_nodes_res()
raise_exception_on_not_ok_status(self._client.lookup_nodes(req, res))
decoder = self.get_node_decoder(node_type)
weights = pywrap.get_node_weights_res(res) \
if decoder.weighted else None
labels = pywrap.get_node_labels_res(res) \
if decoder.labeled else None
int_attrs = pywrap.get_node_int_attr_res(res) \
if decoder.attributed else None
float_attrs = pywrap.get_node_float_attr_res(res) \
if decoder.attributed else None
string_attrs = pywrap.get_node_string_attr_res(res) \
if decoder.attributed else None
int_attrs, float_attrs, string_attrs = \
decoder.format_attrs(int_attrs, float_attrs, string_attrs)
pywrap.del_lookup_nodes_res(res)
pywrap.del_lookup_nodes_req(req)
return Values(int_attrs=int_attrs,
float_attrs=float_attrs,
string_attrs=string_attrs,
weights=weights,
labels=labels,
graph=self)
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 get(self, src_ids, dst_ids):
""" Get batched negative samples.
Args:
src_ids (numpy.array): A 1d numpy array of whose negative dst nodes
will be sampled.
dst_ids (numpy.array): A 1d numpy array of positive dst nodes.
Return:
A `Nodes` object, shape=[ids.shape, `expand_factor`].
"""
src_ids = np.array(src_ids).flatten()
dst_ids = np.array(dst_ids).flatten()
req = self._make_req(src_ids, dst_ids)
res = pywrap.new_sampling_response()
status = self._client.cond_neg_sample(req, res)
if status.ok():
nbrs = pywrap.get_sampling_node_ids(res)
neg_nbrs = self._graph.get_nodes(self._dst_type,
nbrs,
shape=(dst_ids.shape[0],
self._expand_factor))
pywrap.del_op_response(res)
pywrap.del_op_request(req)
errors.raise_exception_on_not_ok_status(status)
return neg_nbrs
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 get(self):
""" Get sampled `SubGraph`.
Return:
An `SubGraph` object.
Raise:
`graphlearn.OutOfRangeError`
"""
state = self._graph.node_state.get(self._seed_type)
req = pywrap.new_subgraph_request(
self._seed_type,
self._nbr_type,
strategy2op(self._strategy, "SubGraphSampler"),
self._batch_size,
state)
res = pywrap.new_subgraph_response()
status = self._client.sample_subgraph(req, res)
if not status.ok():
self._graph.node_state.inc(self._seed_type)
else:
node_ids = pywrap.get_node_set(res)
row_idx = pywrap.get_row_idx(res)
col_idx = pywrap.get_col_idx(res)
edge_ids = pywrap.get_edge_set(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, node_ids)
return SubGraph(np.stack([row_idx, col_idx], axis=0), nodes, Edges(edge_ids=edge_ids))
def lookup_nodes(self, node_type, ids):
""" Get the attributes of given nodes.
"""
ids = np.array(ids)
req = pywrap.new_lookup_nodes_request(node_type)
pywrap.set_lookup_nodes_request(req, ids)
res = pywrap.new_lookup_nodes_response()
status = self._client.lookup_nodes(req, res)
if status.ok():
decoder = self.get_node_decoder(node_type)
weights = pywrap.get_node_weights(
res) if decoder.weighted else None
labels = pywrap.get_node_labels(res) if decoder.labeled else None
int_attrs = pywrap.get_node_int_attributes(
res) if decoder.attributed else None
float_attrs = pywrap.get_node_float_attributes(
res) if decoder.attributed else None
string_attrs = pywrap.get_node_string_attributes(
res) if decoder.attributed else None
int_attrs, float_attrs, string_attrs = decoder.format_attrs(
int_attrs, float_attrs, string_attrs)
pywrap.del_op_response(res)
pywrap.del_op_request(req)
errors.raise_exception_on_not_ok_status(status)
return data.Values(int_attrs=int_attrs,
float_attrs=float_attrs,
string_attrs=string_attrs,
weights=weights,
labels=labels,
graph=self)
def get(self, ids):
""" Get batched samples.
Args:
ids (numpy.array): A 1d numpy array of whose negative dst nodes
will be sampled.
Return:
A `Nodes` object, shape=[ids.shape, `expand_factor`].
"""
if not isinstance(ids, np.ndarray):
raise ValueError("ids must be a numpy array, got {}.".format(
type(ids)))
ids = ids.flatten()
req = self._make_req(ids)
res = pywrap.new_sampling_response()
status = self._client.sample_neighbor(req, res)
if status.ok():
nbrs = pywrap.get_sampling_node_ids(res)
neg_nbrs = self._graph.get_nodes(self._dst_type,
nbrs,
shape=(ids.shape[0],
self._expand_factor))
pywrap.del_op_response(res)
pywrap.del_op_request(req)
raise_exception_on_not_ok_status(status)
return neg_nbrs
def get_entity_count(self, type, is_node):
req = pywrap.new_get_count_request(type, is_node=is_node)
res = pywrap.new_get_count_response()
status = self._client.get_count(req, res)
if status.ok():
count = pywrap.get_count(res)
pywrap.del_op_response(res)
pywrap.del_op_request(req)
errors.raise_exception_on_not_ok_status(status)
return count
def _get_degree(self, edge_type, node_from, ids):
req = pywrap.new_get_degree_request(edge_type, node_from)
pywrap.set_degree_request(req, ids)
res = pywrap.new_get_degree_response()
status = self._client.get_degree(req, res)
if status.ok():
degrees = pywrap.get_degree(res)
pywrap.del_op_response(res)
pywrap.del_op_request(req)
errors.raise_exception_on_not_ok_status(status)
return degrees
def _agg(self, func, segment_ids, num_segments):
req = pywrap.new_agg_nodes_req(self._type,
utils.strategy2op(func, "Aggregator"))
pywrap.set_agg_nodes_req(req, self._ids.flatten(),
np.array(segment_ids, dtype=np.int32),
num_segments)
res = pywrap.new_agg_nodes_res()
raise_exception_on_not_ok_status(self.graph.get_client().agg_nodes(
req, res))
agged = pywrap.get_node_agg_res(res)
pywrap.del_agg_nodes_res(res)
pywrap.del_agg_nodes_req(req)
return agged
def _agg(self, func, segment_ids, num_segments):
req = pywrap.new_aggregating_request(
self._type, utils.strategy2op(func, "Aggregator"))
pywrap.set_aggregating_request(req, self._ids.flatten(),
np.array(segment_ids, dtype=np.int32),
num_segments)
res = pywrap.new_aggregating_response()
status = self.graph.get_client().agg_nodes(req, res)
if status.ok():
agged = pywrap.get_aggregating_nodes(res)
pywrap.del_op_response(res)
pywrap.del_op_request(req)
errors.raise_exception_on_not_ok_status(status)
return agged
def get(self, ids): # pylint: disable=unused-argument
""" Get batched samples.
Args:
ids: A 1d numpy array, the input ids whose neighbors will be returned,
type=np.int64.
Return:
A `Layers` object.
"""
if len(self._meta_path) != len(self._expand_factor):
raise ValueError("The meta_path must have the same number"
"of elements as num_at_each_hop")
src_ids = ids
current_batch_size = ids.size
layers = Layers()
for i in xrange(len(self._meta_path)):
req = self._make_req(i, src_ids)
res = pywrap.new_nbr_res()
status = self._client.sample_neighbor(req, res)
if status.ok():
nbr_ids = pywrap.get_nbr_res_nbr_ids(res)
edge_ids = pywrap.get_nbr_res_edge_ids(res)
pywrap.del_nbr_res(res)
pywrap.del_nbr_req(req)
raise_exception_on_not_ok_status(status)
dst_type = self._dst_types[i]
layer_nodes = self._graph.get_nodes(dst_type,
nbr_ids,
shape=(current_batch_size,
self._expand_factor[i]))
ids = src_ids.repeat(self._expand_factor[i]).flatten()
nbr_ids_flat = nbr_ids.flatten()
layer_edges = \
self._graph.get_edges(self._meta_path[i],
ids,
nbr_ids_flat,
shape=(current_batch_size,
self._expand_factor[i]))
layer_edges.edge_ids = edge_ids
layers.append_layer(Layer(layer_nodes, layer_edges))
current_batch_size = nbr_ids_flat.size
src_ids = nbr_ids
return layers
def get(self, ids): # pylint: disable=unused-argument
if len(self._meta_path) != len(self._expand_factor):
raise ValueError("The length of meta_path must be same with hop count.")
ids = np.array(ids).flatten()
src_ids = ids
current_batch_size = ids.size
layers = Layers()
for i in range(len(self._meta_path)):
# req, res & call method.
req = self._make_req(i, src_ids)
res = pywrap.new_sampling_response()
status = self._client.sample_neighbor(req, res)
if status.ok():
src_degrees = pywrap.get_sampling_node_degrees(res)
dense_shape = (current_batch_size, max(src_degrees))
nbr_ids = pywrap.get_sampling_node_ids(res)
edge_ids = pywrap.get_sampling_edge_ids(res)
pywrap.del_op_response(res)
pywrap.del_op_request(req)
errors.raise_exception_on_not_ok_status(status)
dst_type = self._dst_types[i]
layer_nodes = self._graph.get_nodes(
dst_type,
nbr_ids,
offsets=src_degrees,
shape=dense_shape)
ids = np.concatenate([src_ids[idx].repeat(d) for \
idx, d in enumerate(src_degrees)])
nbr_ids_flat = nbr_ids.flatten()
layer_edges = self._graph.get_edges(
self._meta_path[i],
ids,
nbr_ids_flat,
offsets=src_degrees,
shape=dense_shape)
layer_edges.edge_ids = edge_ids
layers.append_layer(Layer(layer_nodes, layer_edges))
current_batch_size = nbr_ids_flat.size
src_ids = nbr_ids
return layers
def get(self, ids): # pylint: disable=unused-argument
if len(self._meta_path) != len(self._expand_factor):
raise ValueError("The meta_path must have the same number"
"of elements as num_at_each_hop")
ids = ids.flatten()
src_ids = ids
current_batch_size = ids.size
layers = Layers()
for i in xrange(len(self._meta_path)):
# req, res & call method.
req = self._make_req(i, src_ids)
res = pywrap.new_nbr_res()
raise_exception_on_not_ok_status(
self._client.sample_neighbor(req, res))
src_degrees = pywrap.get_nbr_res_degrees(res)
dense_shape = (current_batch_size, max(src_degrees))
nbr_ids = pywrap.get_nbr_res_nbr_ids(res)
edge_ids = pywrap.get_nbr_res_edge_ids(res)
pywrap.del_nbr_res(res)
pywrap.del_nbr_req(req)
dst_type = self._dst_types[i]
layer_nodes = self._graph.get_nodes(dst_type,
nbr_ids,
offsets=src_degrees,
shape=dense_shape)
ids = np.concatenate([src_ids[idx].repeat(d) for \
idx, d in enumerate(src_degrees)])
nbr_ids_flat = nbr_ids.flatten()
layer_edges = \
self._graph.get_edges(self._meta_path[i],
ids,
nbr_ids_flat,
offsets=src_degrees,
shape=dense_shape)
layer_edges.edge_ids = edge_ids
layers.append_layer(Layer(layer_nodes, layer_edges))
current_batch_size = nbr_ids_flat.size
src_ids = nbr_ids
return layers
def __init__(self, dag, capacity=10):
assert dag.is_ready(), \
"Query should start with E()/V() and end with value()."
assert isinstance(capacity, int) and 0 < capacity < 128, \
"Dataset capacity should in range of (0, 128)."
self._dag = dag
self._dag_id = dag.name
self._cur_res = None
graph = dag.graph
client = graph.get_client()
status = client.run_dag(self._dag.dag_def)
raise_exception_on_not_ok_status(status)
pywrap.set_dataset_capacity(capacity)
self._dag_dataset = pywrap.Dataset(client, self._dag_id)
graph.add_dataset(self)
def lookup_edges(self, edge_type, src_ids, edge_ids):
""" Get the attributes of given edges.
"""
src_ids = np.array(src_ids)
edge_ids = np.array(edge_ids)
batch_size = src_ids.flatten().size
if batch_size != edge_ids.flatten().size:
raise ValueError("src_ids and edge_ids for lookup edges must "
"be same, got {} and {}".format(
batch_size,
edge_ids.flatten().size))
req = pywrap.new_lookup_edges_request(edge_type)
pywrap.set_lookup_edges_request(req, src_ids, edge_ids)
res = pywrap.new_lookup_edges_response()
status = self._client.lookup_edges(req, res)
if status.ok():
decoder = self.get_edge_decoder(edge_type)
weights = pywrap.get_edge_weights(
res) if decoder.weighted else None
labels = pywrap.get_edge_labels(res) if decoder.labeled else None
int_attrs = pywrap.get_edge_int_attributes(
res) if decoder.attributed else None
float_attrs = pywrap.get_edge_float_attributes(
res) if decoder.attributed else None
string_attrs = pywrap.get_edge_string_attributes(
res) if decoder.attributed else None
int_attrs, float_attrs, string_attrs = decoder.format_attrs(
int_attrs, float_attrs, string_attrs)
pywrap.del_op_response(res)
pywrap.del_op_request(req)
errors.raise_exception_on_not_ok_status(status)
return data.Values(int_attrs=int_attrs,
float_attrs=float_attrs,
string_attrs=string_attrs,
weights=weights,
labels=labels,
graph=self)
def lookup_edges(self, edge_type, src_ids, edge_ids):
""" Get all the edge properties.
"""
batch_size = src_ids.flatten().size
if batch_size != edge_ids.flatten().size:
raise ValueError("src_ids and edge_ids for lookup edges must "
"be same, got {} and {}".format(
batch_size,
edge_ids.flatten().size))
req = pywrap.new_lookup_edges_req(edge_type)
pywrap.set_lookup_edges_req(req, src_ids, edge_ids)
res = pywrap.new_lookup_edges_res()
status = self._client.lookup_edges(req, res)
if status.ok():
decoder = self.get_edge_decoder(edge_type)
weights = pywrap.get_edge_weights_res(res) \
if decoder.weighted else None
labels = pywrap.get_edge_labels_res(res) \
if decoder.labeled else None
int_attrs = pywrap.get_edge_int_attr_res(res) \
if decoder.attributed else None
float_attrs = pywrap.get_edge_float_attr_res(res) \
if decoder.attributed else None
string_attrs = pywrap.get_edge_string_attr_res(res) \
if decoder.attributed else None
int_attrs, float_attrs, string_attrs = \
decoder.format_attrs(int_attrs, float_attrs, string_attrs)
pywrap.del_lookup_edges_res(res)
pywrap.del_lookup_edges_req(req)
raise_exception_on_not_ok_status(status)
return Values(int_attrs=int_attrs,
float_attrs=float_attrs,
string_attrs=string_attrs,
weights=weights,
labels=labels,
graph=self)
def get(self):
""" Get batched sampled `Edges`.
Return:
An `Edges` object, shape=[`batch_size`]
Raise:
`graphlearn.OutOfRangeError`
"""
req = pywrap.new_get_edge_req(self._edge_type,
self._strategy,
self._batch_size)
res = pywrap.new_get_edge_res()
raise_exception_on_not_ok_status(self._client.get_edges(req, res))
src_ids = pywrap.get_edge_src_id_res(res)
dst_ids = pywrap.get_edge_dst_id_res(res)
edge_ids = pywrap.get_edge_edge_id_res(res)
edges = self._graph.get_edges(self._edge_type,
src_ids,
dst_ids)
edges.edge_ids = edge_ids
return edges
