Solve bug in sampler when the same node is visited twice from different path in the sampling plan.

dgf/src/sampling/in_memory_sampler_ext.cc

@@ -29,6 +29,7 @@
 #include "absl/container/btree_map.h"
 #include "absl/container/btree_set.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -215,6 +216,10 @@ struct SampleBuilder {
   // depth.
   std::vector<std::vector<InputIdx>> recursion_cache;
 
+  // Keep track of visited nodes per plan step to avoid multi-visits on the same
+  // step.
+  std::vector<absl::flat_hash_set<InputIdx>> visited_node_idxs;
+
   // Random number generator.
   std::mt19937_64 rng;
 
@@ -300,9 +305,17 @@ struct SampleBuilder {
                 target_node);
             DCHECK(sample_target_nodeset.sampled_node_idx_to_node_idx.size() ==
                    sample_target_nodeset.node_idx_to_sampled_node_idx.size());
-          } else if (!sampler.plan_.multi_visit) {
-            recuse = false;
           }
+
+          if (!sampler.plan_.multi_visit) {
+            // Further expand the node iff it was not already visited in
+            // this plan edge. Note that the same node can be expanded multiple
+            // times through different plan edges.
+            auto [_it, inserted_local] =
+                visited_node_idxs[plan_edge.node->step_idx].insert(target_node);
+            recuse = inserted_local;
+          }
+
         } else {
           // Sampling with replacement.
           target_sampled_node =
@@ -330,6 +343,14 @@ struct SampleBuilder {
     if (recursion_cache.size() < sampler.plan_.num_steps) {
       recursion_cache.resize(sampler.plan_.num_steps);
     }
+    if (!sampler.plan_.multi_visit) {
+      if (visited_node_idxs.size() < sampler.plan_.num_steps) {
+        visited_node_idxs.resize(sampler.plan_.num_steps);
+      }
+      for (auto& s : visited_node_idxs) {
+        s.clear();
+      }
+    }
 
     // Reuse capacity of edgesets and nodesets.
     if (edgesets.size() < sampler.edgesets_.size()) {
@@ -1112,7 +1133,7 @@ absl::Status Sampler::IndexEdgeSets(const nb::object& py_graph,
   // Index the nodesets.
   nodesets_.assign(schema_->nodeset_name_to_idx.size(), {});
   DGF_GET_ATTR_OR_RETURN(nb::dict, py_node_sets, py_graph, "node_sets");
-  for (const auto nodeset : schema_->nodeset_name_to_idx) {
+  for (const auto& nodeset : schema_->nodeset_name_to_idx) {
     DGF_ASSIGN_OR_RETURN(const nb::object py_nodeset,
                          GetItemFromPyDict<nb::object>(
                              py_node_sets, string_to_py_str(nodeset.first)));

dgf/src/sampling/in_memory_sampler_test.py

@@ -1525,9 +1525,7 @@ def test_slice_timeseries_by_seed_validation_errors(self):
         " `slice_timeseries_by_seed=True` and the schema contains"
         " `is_timeseries=True` features.",
     ):
-      sampler_missing_timestamps.sample(
-          seed_node_idxs=0, seed_timestamps=None
-      )
+      sampler_missing_timestamps.sample(seed_node_idxs=0, seed_timestamps=None)
 
   def test_timeseries_subgraph_and_multisubgraph_clipping(self):
     graph, schema = self._create_causal_timeseries_test_graph()
@@ -1672,14 +1670,15 @@ def setUpClass(cls):
         },
     )
 
-  def test_sample(self):
+  @parameterized.parameters(True, False)
+  def test_multi_visit(self, multi_visit: bool):
     plan = config_lib.simple_sampling_config_to_sampling_plan(
         config_lib.SimpleSamplingConfig(
             seed_nodeset="n1",
             num_hops=5,
             hop_width=2,
             reverse=False,
-            multi_visit=False,
+            multi_visit=multi_visit,
         ),
         self.schema,
     )
@@ -1693,8 +1692,10 @@ def test_sample(self):
     )
     sample = sampler.sample(0)
     self.assertIsNotNone(sample)
-    # The diamon and exactly 2 or the last 3 nodes are sampled.
-    self.assertEqual(sample.node_sets["n1"].num_nodes, 6)
+    if multi_visit:
+      self.assertEqual(sample.node_sets["n1"].num_nodes, 7)
+    else:
+      self.assertEqual(sample.node_sets["n1"].num_nodes, 6)
 
 
 if __name__ == "__main__":