support standard temporal sampling and updated PLC parameters

Author: alexbarghi-nv

python/cugraph-pyg/cugraph_pyg/loader/link_neighbor_loader.py

@@ -59,6 +59,7 @@ def __init__(
         batch_size: int = 16,  # Refers to number of edges per batch.
         compression: Optional[str] = None,
         local_seeds_per_call: Optional[int] = None,
+        temporal_comparison: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -128,12 +129,19 @@ def __init__(
             all workers.  If not provided, it will be automatically
             calculated.
             See cugraph_pyg.sampler.BaseDistributedSampler.
+        temporal_comparison: str (optional, default='<=', i.e. monotonically decreasing)
+            The comparison operator for temporal sampling (>, <, >=, <=, last).
+            Note that this should be 'last' for temporal_strategy='last'.
+            See cugraph_pyg.sampler.BaseDistributedSampler.
         **kwargs
             Other keyword arguments passed to the superclass.
         """
 
         subgraph_type = torch_geometric.sampler.base.SubgraphType(subgraph_type)
 
+        if temporal_comparison is None:
+            temporal_comparison = "<="
+
         if not directed:
             subgraph_type = torch_geometric.sampler.base.SubgraphType.induced
             warnings.warn(
@@ -172,17 +180,15 @@ def __init__(
 
         is_temporal = (edge_label_time is not None) and (time_attr is not None)
 
+        if (edge_label_time is None) != (time_attr is None):
+            warnings.warn(
+                "Edge-based temporal sampling requires that both edge_label_time and time_attr are provided. Defaulting to non-temporal sampling."
+            )
+
         if weight_attr is not None:
             graph_store._set_weight_attr((feature_store, weight_attr))
         if is_temporal:
-            # TODO Confirm that time is an edge attribute
-            # TODO Add support for time override (see rapidsai/cugraph#5263)
             graph_store._set_etime_attr((feature_store, time_attr))
-            warnings.warn(
-                "Temporal sampling in cuGraph-PyG is currently only forward in time"
-                " instead of the expected backward in time.  This will be fixed in a"
-                " future release."
-            )
 
         if isinstance(num_neighbors, dict):
             sorted_keys, _, _ = graph_store._numeric_edge_types
@@ -209,6 +215,7 @@ def __init__(
                 biased=(weight_attr is not None),
                 heterogeneous=(not graph_store.is_homogeneous),
                 temporal=is_temporal,
+                temporal_comparison=temporal_comparison,
                 vertex_type_offsets=graph_store._vertex_offset_array,
                 num_edge_types=len(graph_store.get_all_edge_attrs()),
             ),

python/cugraph-pyg/cugraph_pyg/loader/neighbor_loader.py

@@ -56,6 +56,7 @@ def __init__(
         batch_size: int = 16,
         compression: Optional[str] = None,
         local_seeds_per_call: Optional[int] = None,
+        temporal_comparison: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -121,12 +122,19 @@ def __init__(
             all workers.  If not provided, it will be automatically
             calculated.
             See cugraph_pyg.sampler.BaseDistributedSampler.
+        temporal_comparison: str (optional, default='<=', i.e. monotonically decreasing)
+            The comparison operator for temporal sampling (>, <, >=, <=, last).
+            Note that this should be 'last' for temporal_strategy='last'.
+            See cugraph_pyg.sampler.BaseDistributedSampler.
         **kwargs
             Other keyword arguments passed to the superclass.
         """
 
         subgraph_type = torch_geometric.sampler.base.SubgraphType(subgraph_type)
 
+        if temporal_comparison is None:
+            temporal_comparison = "<="
+
         if not directed:
             subgraph_type = torch_geometric.sampler.base.SubgraphType.induced
             warnings.warn(
@@ -166,15 +174,18 @@ def __init__(
         is_temporal = time_attr is not None
 
         if is_temporal:
-            # TODO Confirm that time is an edge attribute
-            # TODO Add support for time override (see rapidsai/cugraph#5263)
             graph_store._set_etime_attr((feature_store, time_attr))
 
-            warnings.warn(
-                "Temporal sampling in cuGraph-PyG is currently only forward in time"
-                " instead of the expected backward in time.  This will be fixed in a"
-                " future release."
-            )
+            if input_time is None:
+                input_type, input_nodes, _ = (
+                    torch_geometric.loader.utils.get_input_nodes(
+                        data, input_nodes, None
+                    )
+                )
+                if input_type is None:
+                    input_type = list(graph_store._vertex_offsets.keys())[0]
+                # will assume the time attribute exists for nodes as well
+                input_time = feature_store[input_type, time_attr, None][input_nodes]
 
         if weight_attr is not None:
             graph_store._set_weight_attr((feature_store, weight_attr))
@@ -204,6 +215,7 @@ def __init__(
                 biased=(weight_attr is not None),
                 heterogeneous=(not graph_store.is_homogeneous),
                 temporal=is_temporal,
+                temporal_comparison=temporal_comparison,
                 vertex_type_offsets=graph_store._vertex_offset_array,
                 num_edge_types=len(graph_store.get_all_edge_attrs()),
             ),

python/cugraph-pyg/cugraph_pyg/loader/node_loader.py

@@ -116,7 +116,7 @@ def __init__(
             if input_id is None
             else input_id,
             node=input_nodes,
-            time=None,
+            time=input_time,
             input_type=input_type,
         )