dgf.sampling

dgf.sampling.Sampler

Sampler for generating subgraphs from an in-memory graph.

`multisubgraph`

Extracts the subgraphs around the provided seed nodes.

This method returns the graphs containing all the nodes and edges at a distance less than or equal to the configured number of hops from the provided seed_node_idxs. Each seed-node leads to the creation of a different sub-graph independently.

The seed nodes are always the first node of the extracted graph.

Functionally, multisubgraph returns the same results as sample with an infinite width, but is massively more efficient. Both multisubgraph and sample use a graph traversal algorithm. However, multisubgraph doesn't re-visit nodes, which can make it more efficient.

Usage example:

graph, schema = dgf.io.read_graph(<path to graph>)
config = dgf.sampling.SimpleSamplingConfig(
    seed_nodeset="client",
    num_hops=4,
    hop_width=1, # Not used with "subgraph".
    reverse=True,
)
sampler = dgf.sampling.create_sampler(graph, config, schema)

subgraph = sampler.subgraph([0,1,2])
print(subgraph)

Parameters:

Name	Type	Description	Default
`seed_node_idxs`	`List[int]`	The indexes of the nodes to start sampling from.	required

Returns:

Type	Description
`List[InMemoryGraph]`	A list of graphs. One for each "seed_node_idxs" value.

`sample`

sample

sample

Samples one (or multiple) subgraphs.

Grows one or more graph samples starting from the provided seed nodes. Each graph sample is constructed by randomly traversing edges and aggregating all visited edges and nodes.

Parameters:

Name	Type	Description	Default
`seed_node_idxs`	`Union[int, List[int], ndarray]`	The index or indexes of the nodes to start sampling from. Sampling multiple nodes at the same time is more efficient that calling "sample" multiple times.	required
`seed_timestamps`	`Optional[Union[int, List[int], ndarray]]`	Optional timestamps for time-aware sampling. If specified, only sample edges with a timestamp anterior (non strict) to the provided seed_timestamp. Should have the same length as "seed_node_idxs". Requires for the sampler to be initialized with some timestamps.	`None`
`masked_edge_idxs`	`Optional[Union[int, List[int], ndarray]]`	Optional edge indices to mask during sampling. If specified, masks the specified edge. Should have the same length as "seed_node_idxs". Requires for the sampler to be initialized with a masked edgeset. If a seed node idx is -1, not edge filtering is done.	`None`

Returns:

Type	Description
`Union[InMemoryGraph, List[InMemoryGraph]]`	An `InMemoryGraph` or list of `InMemoryGraph`
`Union[InMemoryGraph, List[InMemoryGraph]]`	representing the sampled subgraph.

`set_return_options`

Sets whether to return features and node indices in sampled graphs.

Parameters:

Name	Type	Description	Default
`return_features`	`bool`	Whether to include feature values in the returned graph.	required
`return_node_idxs`	`bool`	Whether to include node indexes in the returned graph as a "#idx" node feature.	required

`subgraph`

Extracts the subgraph around the provided seed nodes.

This method returns a graph containing all the nodes and edges at a distance less than or equal to the configured number of hops from the provided seed_node_idxs.

The seed nodes are always the first nodes of their respective node set in the returned InMemoryGraph. For example, if seed_node_idxs contains 3 elements from a specific node set, these will correspond to the first 3 nodes of that same node set in the extracted graph.

Warning: Unlike "sample" that returns a different graph for each seed-node, "subgraph" returns a single possibly connected graph. To compute independent subgraphs, use multisubgraph instead.

Usage example:

graph, schema = dgf.io.read_graph(<path to graph>)
config = dgf.sampling.SimpleSamplingConfig(
    seed_nodeset="client",
    num_hops=4,
    hop_width=1, # Not used with "subgraph".
    reverse=True,
)
sampler = dgf.sampling.create_sampler(graph, config, schema)

subgraph = sampler.subgraph([0,1,2])
print(subgraph)

Parameters:

Name	Type	Description	Default
`seed_node_idxs`	`List[int]`	The indexes of the nodes to start sampling from.	required

Returns:

Type	Description
`InMemoryGraph`	The resulting graph.

dgf.sampling.SamplingPlan

Defines a complex sampling config.

Attributes:

Name	Type	Description
`root`	`PlanNode`	The root node of the sampling plan, specifying the starting nodeset.
`with_replacement`	`bool`	Test if the sampling is done with replacement. See documentation for "with_replacement" attribute in SimpleSamplingConfig for the full explanation.
`temporal_sampling`	`bool`	If True, temporal sampling is enabled and causal timestamps are inferred from the schema.
`edgeset_timestamp_features`	`Dict[str, str]`	Mapping from edgeset name to its timestamp feature name for causal filtering.

dgf.sampling.SimpleSamplingConfig

Configuration for simple neighborhood sampling.

This configuration defines a simple, breadth-first sampling strategy. Starting from a specified seed_nodeset, it performs a fixed number of num_hops. At each hop, for every edge type connected to the current nodeset, up to hop_width neighbors are sampled.

Attributes:

Name	Type	Description
`seed_nodeset`	`str`	The name of the nodeset from which to start the sampling.
`num_hops`	`int`	The maximum number of hops (steps) to perform outwards from the `seed_nodeset`. This determines the depth of the sampled neighborhood.
`hop_width`	`int`	The maximum number of neighbors to sample for each edge type at every hop. If more than `hop_width` neighbors are available, a subset is chosen.
`reverse`	`bool`	If True, edges can be traversed in both their defined direction and the reverse direction. If False, only the defined forward direction of edges is used.
`with_replacement`	`bool`	If false, the sampled graph is a sub-graph of the original graph with cycles (if the original graph has cycles). Nodes / edges that are visited multiple times do not lead to multiple nodes / edges in the sampled graph. If true, the sampled graph is a tree where nodes / edges in the original graph might lead to multiple nodes / edges in the sampled grpah.
`temporal_sampling`	`bool`	If True, temporal sampling is enabled and causal timestamps are inferred from the schema.

dgf.sampling.SpannerGraphSampler

Sampler that executes queries on Spanner directly to fetch subgraphs.

`sample`

Samples subgraphs starting from the given seed nodes.

Parameters:

Name	Type	Description	Default
`seed_ids`	`List[bytes]`	The list of node IDs (as bytes) to use as seeds.	required

Returns:

Type	Description
`List[InMemoryGraph]`	A list of InMemoryGraph objects corresponding to the subgraphs sampled for
`List[InMemoryGraph]`	each seed ID.

dgf.sampling.create_graph_spanner_sampler

Creates a SpannerGraphSampler instance.

dgf.sampling.create_sampler

Creates an in-memory sampler.

Parameters:

Name	Type	Description	Default
`graph`	`InMemoryGraph`	The in-memory heterogeneous graph to sample from.	required
`plan`	`Union[SimpleSamplingConfig, SamplingPlan]`	The sampling plan configuration. Can be a `SimpleSamplingConfig` or a `SamplingPlan`.	required
`schema`	`GraphSchema`	Graph schema. Required if `plan` is a `SimpleSamplingConfig`.	required
`batch_size`	`Optional[int]`	Number of samples you will sample each time. Sampling more / less samples at the same time is possible but possibly less efficient..	`None`
`return_features`	`bool`	Whether to include feature values in the returned graph.	`True`
`return_node_idxs`	`bool`	Whether to include node indexes in the returned graph in as a "#idx" node feature.	`False`
`debug_sampling`	`bool`	If true, enables a deterministic debug mode. In this mode, sampling always selects the first available edges, making the process fully reproducible.	`False`
`num_threads`	`Optional[int]`	Number of sampling threads. If None, select the number of threads automatically. Set num_threads=0 to disable multi-threading.	`None`
`seed`	`Optional[int]`	A positive integer to use as random seed for the sampler. If not provided, the seed is randomly initialized. Note that variation in the compilation can lead to variation (e.g., recompiling the binary might lead to different results--though this should be rare). Note: For writing unit tests,using debug_sampling=True is better.	`None`
`edgeset_to_mask`	`Optional[str]`	Optional edgeset name to mask during sampling.	`None`
`slice_timeseries_by_seed`	`Optional[bool]`	Whether to causally slice `is_timeseries=True` sequence features by the seed node timestamp. Defaults to `plan.temporal_sampling`.	`None`
`max_timeseries_len`	`Optional[int]`	Optional cap on the number of historical causal sequence steps retained for each timeseries feature.	`None`

TODO(gbm): Should we remove the compilation variations (e.g., change in random number generator, change in hashmaps).

Returns:

Type	Description
`Sampler`	A `Sampler` instance.

dgf.sampling.extract_beam_nodes_ids

Extracts all the node ids of a given nodeset.

This method can be used to create the seed nodes argument for the sampler.

See "sample_with_beam_semi_distributed_sampler" for a usage example.

Parameters:

Name	Type	Description	Default
`graph`	`Graph`	The distributed graph containing the node sets.	required
`target_nodeset`	`str`	The name of the node set from which to extract node IDs.	required

Returns:

Type	Description
`PCollection[NodeId]`	A `PCollection` of bytes, where each element is a unique node ID
`PCollection[NodeId]`	from the specified `target_nodeset`.

dgf.sampling.sample_with_beam_semi_distributed_sampler

Samples subgraphs from a distributed graph using a semi-distributed algo.

WARNING: This sampler is not operational. Use the "sample_with_beam_semi_distributed_sampler_v2".

This beam sampler generates samples by running the in-process sampler multiple times in parallel on different workers. Only the final feature gathering is distributed. This sampler is suited for graph where the topologie fits in memory. For reference, a graph with 1B edges and 100M nodes takes (1B + 100M) * 8 = ~9GB of RAM (assuming uint64 indexing, no compression).

Usage example:

# Read HGraph
graph = dgf.io.ReadFromHGraph("/cns/.../my_hgraph")

# Create the sampling config / plan
sampling_config = dgf.sampling.SimpleSamplingConfig(
    seed_nodeset="paper", num_hops=2, hop_width=3)
sampling_plan = dgf.sampling.simple_sampling_config_to_sampling_plan(
    sampling_config,
    graph.schema,
)

# Get all the nodes as seeds
seeds = dgf.sampler.extract_beam_nodes_ids(graph, "paper")

# Generate samples
samples = dgf.sampler.sample_with_beam_semi_distributed_sampler(
    graph, plan seeds)

# Save the samples to disk
dgf.io.write_to_tf_graph_sample(samples, "/cns/.../samples@*")

Parameters:

Name	Type	Description	Default
`graph`	`Graph`	The distributed graph to sample from..	required
`plan`	`SamplingPlan`	The sampling plan..	required
`seeds`	`PCollection[NodeId]`	PCollection of node IDs to use as seed. These IDs must belong to the nodeset specified in `plan.root.nodeset`.	required
`debug_sampling`	`bool`	If true, enables debug mode in the sampler. Used for unit testing.	`False`

Returns:

Type	Description
`PCollection[PKeyedInMemoryGraph]`	A `PCollection` of `InMemoryHeterogeneousGraph` instances, where each
`PCollection[PKeyedInMemoryGraph]`	instance represents a sampled subgraph.

dgf.sampling.sample_with_beam_semi_distributed_sampler_v2

Samples subgraphs from a distributed graph using a semi-distributed algo.

This beam sampler generates samples by running the in-process sampler multiple times in parallel on different workers. Only the final feature gathering is distributed. This sampler is suited for graph where the topology fits in memory. For reference, a graph with 1B edges and 100M nodes takes (1B + 100M) * 8 = ~9GB of RAM (assuming uint64 indexing, no compression).

Usage example:

graph_path = ...

# Use all the nodes as seed nodes.
graph = dgf.beam.io.read_graph(root, graph_path)
seeds = dgf.sampler.extract_beam_nodes_ids(graph, "paper")
seeds = seeds | "Reshuffle seeds" >> beam.Reshuffle()

# Create the sampling config
sampling_config = dgf.sampling.SimpleSamplingConfig(
    seed_nodeset="paper", num_hops=2, hop_width=3)

# Generate samples
samples = dgf.sampler.sample_with_beam_semi_distributed_sampler_v2(
    graph_path, sampling_config, seeds)

# Save the samples to disk
dgf.io.write_to_tf_graph_sample(samples, "/cns/.../samples@*")

Parameters:

Name	Type	Description	Default
`graph_path`	`str`	Path to graph on disk.	required
`plan`	`Union[SimpleSamplingConfig, SamplingPlan]`	The sampling plan..	required
`seeds`	`PCollection[NodeId]`	PCollection of node IDs to use as seed. These IDs must belong to the nodeset specified in `plan.root.nodeset`.	required
`debug_sampling`	`bool`	If true, enables debug mode in the sampler. Used for unit testing.	`False`
`num_threads`	`int`	The number of threads used by each worker.	`20`
`beam_feature_collection`	`bool`	If False (default), the feature values are gathered by the in-memory sampler. This option is fast but requires more RAM. If false, the feature values are gathered by a Beam join after the sampling.	`False`
`beam_namespace`	`str`	Prefix added to all the beam ptransforms.	`''`

Returns:

Type	Description
`PKeyedInMemoryGraph`	A `PCollection` of `InMemoryHeterogeneousGraph` instances, where each
`GraphSchema`	instance represents a sampled subgraph.

dgf.sampling.simple_sampling_config_to_sampling_plan

Converts a SimpleSamplingConfig to a more general SamplingPlan.

Parameters:

Name	Type	Description	Default
`src`	`SimpleSamplingConfig`	The SimpleSamplingConfig to convert.	required
`schema`	`GraphSchema`	The graph schema, used to resolve edge connections.	required

Returns:

Type	Description
`SamplingPlan`	A SamplingPlan equivalent to the provided SimpleSamplingConfig.