Preserve temporal sampling cutoffs with cuGraph 26.10

Signed-off-by: Ard <88337265+aw471@users.noreply.github.com>

Author: aw471

sdm/relational/backend/_cugraph.py

@@ -196,7 +196,10 @@ def _build_graph(self) -> None:
         dst = torch.cat(dsts)
         self._num_edges = src.numel()
         edge_type = torch.cat(edge_types)
-        edge_time = torch.cat(edge_times) if self.time_columns else None
+        # Reverse time so the task cutoff is a lower bound, including NaT.
+        edge_time = (
+            torch.cat(edge_times).bitwise_not() if self.time_columns else None
+        )
         self._has_outgoing = torch.zeros(
             self._num_vertices,
             dtype=torch.bool,
@@ -484,13 +487,14 @@ def _sample_temporal_hop(
         for table_name in possible_source_tables:
             fanout[self._outgoing_edge_types[table_name]] = count
 
+        # Temporal disjoint sampling uses cuGraph's version-specific default.
         result = (
             self._pylibcugraph.heterogeneous_uniform_temporal_neighbor_sample(
                 self._resource_handle,
                 self._graph,
                 "edge_start_time",
                 cp.from_dlpack(frontier_node),
-                cp.from_dlpack(seed_time[frontier_example]),
+                cp.from_dlpack(seed_time[frontier_example].bitwise_not()),
                 cp.from_dlpack(label_offsets),
                 cp.from_dlpack(self._vertex_offsets),
                 fanout,
@@ -499,14 +503,13 @@ def _sample_temporal_hop(
                 do_expensive_check=False,
                 prior_sources_behavior=None,
                 deduplicate_sources=True,
-                disjoint_sampling=False,
                 return_hops=False,
                 renumber=False,
                 retain_seeds=False,
                 compression="COO",
                 compress_per_hop=False,
                 random_state=self._next_random_state(),
-                temporal_sampling_comparison="monotonically_decreasing",
+                temporal_sampling_comparison="monotonically_increasing",
             )
         )
         minor = self._as_tensor(result["minors"])

test/relational/backend/test_cugraph.py

@@ -385,15 +385,17 @@ def capture(*args: Any, **kwargs: Any) -> Any:
 
 
 @onlyCUDA
+@pytest.mark.parametrize("num_neighbors", [[-1, -1], [3, 3]])
 def test_cugraph_sampler_uses_original_cutoff(
     temporal_data: RelationalData,
+    num_neighbors: list[int],
 ) -> None:
     _require_rapids()
     data = temporal_data.cuda()
     task_table = _table(
         {
-            "entity": [0, 0],
-            "cutoff": pd.to_datetime([2, 10], unit="s").tolist(),
+            "entity": [0, 0, 0],
+            "cutoff": pd.to_datetime([0, 2, 10], unit="s").tolist(),
         },
         {"entity": Stype.id, "cutoff": Stype.datetime},
     )
@@ -405,19 +407,22 @@ def test_cugraph_sampler_uses_original_cutoff(
             "table": "roots",
             "table_column": "root_id",
         },
-        num_neighbors=[-1, -1],
+        num_neighbors=num_neighbors,
         task_time_column="cutoff",
         temporal_strategy="uniform",
     )
 
     assert _rows(
         output.related_tables.tables["first"], EXAMPLE_ID, "first_id"
-    ) == [(0, 11), (0, 12), (1, 10), (1, 11), (1, 12)]
+    ) == [(1, 11), (1, 12), (2, 10), (2, 11), (2, 12)]
     # The second-hop time (8) is newer than the first-hop time (3), but is
     # valid because every hop uses the task cutoff (10).
     assert _rows(
         output.related_tables.tables["second"], EXAMPLE_ID, "second_id"
-    ) == [(1, 20), (1, 21)]
+    ) == [(2, 20), (2, 21)]
+    assert _rows(
+        output.related_tables.tables["roots"], EXAMPLE_ID, "root_id"
+    ) == [(0, 0), (1, 0), (2, 0)]
 
 
 @onlyCUDA