[GraphBolt] Add `__repr__` to feature store. (#6788)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-0-133.us-west-2.compute.internal>

python/dgl/graphbolt/impl/torch_based_feature_store.py

@@ -168,6 +168,9 @@ class TorchBasedFeature(Feature):
         """In-place operation to copy the feature to pinned memory."""
         self._tensor = self._tensor.pin_memory()
 
+    def __repr__(self) -> str:
+        return _torch_based_feature_str(self)
+
 
 class TorchBasedFeatureStore(BasicFeatureStore):
     r"""A store to manage multiple pytorch based feature for access.
@@ -231,3 +234,42 @@ class TorchBasedFeatureStore(BasicFeatureStore):
         """In-place operation to copy the feature store to pinned memory."""
         for feature in self._features.values():
             feature.pin_memory_()
+
+    def __repr__(self) -> str:
+        return _torch_based_feature_store_str(self._features)
+
+
+def _torch_based_feature_str(feature: TorchBasedFeature) -> str:
+    final_str = "TorchBasedFeature("
+    indent_len = len(final_str)
+
+    def _add_indent(_str, indent):
+        lines = _str.split("\n")
+        lines = [lines[0]] + [" " * indent + line for line in lines[1:]]
+        return "\n".join(lines)
+
+    feature_str = "feature=" + _add_indent(
+        str(feature._tensor), indent_len + len("feature=")
+    )
+    final_str += feature_str + ",\n" + " " * indent_len
+    metadata_str = "metadata=" + _add_indent(
+        str(feature.metadata()), indent_len + len("metadata=")
+    )
+    final_str += metadata_str + ",\n)"
+    return final_str
+
+
+def _torch_based_feature_store_str(
+    features: Dict[str, TorchBasedFeature]
+) -> str:
+    final_str = "TorchBasedFeatureStore"
+    indent_len = len(final_str)
+
+    def _add_indent(_str, indent):
+        lines = _str.split("\n")
+        lines = [lines[0]] + [" " * indent + line for line in lines[1:]]
+        return "\n".join(lines)
+
+    features_str = _add_indent(str(features), indent_len)
+    final_str += features_str
+    return final_str

tests/python/pytorch/graphbolt/impl/test_torch_based_feature_store.py

@@ -275,3 +275,85 @@ def test_torch_based_feature_store(in_memory):
         assert feature_store.size("node", None, "a") == torch.Size([3])
 
         feature_store = None
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_torch_based_feature_repr(in_memory):
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 3], [4, 5, 6]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
+        metadata = {"max_value": 3}
+        if not in_memory:
+            a = to_on_disk_tensor(test_dir, "a", a)
+            b = to_on_disk_tensor(test_dir, "b", b)
+
+        feature_a = gb.TorchBasedFeature(a, metadata=metadata)
+        feature_b = gb.TorchBasedFeature(b)
+
+        expected_str_feature_a = str(
+            """TorchBasedFeature(feature=tensor([[1, 2, 3],
+                                  [4, 5, 6]]),
+                  metadata={'max_value': 3},
+)"""
+        )
+        expected_str_feature_b = str(
+            """TorchBasedFeature(feature=tensor([[[1, 2],
+                                   [3, 4]],
+                          
+                                  [[4, 5],
+                                   [6, 7]]]),
+                  metadata={},
+)"""
+        )
+        assert str(feature_a) == expected_str_feature_a
+        assert str(feature_b) == expected_str_feature_b
+        a = b = metadata = None
+        feature_a = feature_b = None
+        expected_str_feature_a = expected_str_feature_b = None
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_torch_based_feature_store_repr(in_memory):
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 4], [2, 5, 3]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
+        write_tensor_to_disk(test_dir, "a", a, fmt="torch")
+        write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
+        feature_data = [
+            gb.OnDiskFeatureData(
+                domain="node",
+                type="paper",
+                name="a",
+                format="torch",
+                path=os.path.join(test_dir, "a.pt"),
+                in_memory=True,
+            ),
+            gb.OnDiskFeatureData(
+                domain="edge",
+                type="paper:cites:paper",
+                name="b",
+                format="numpy",
+                path=os.path.join(test_dir, "b.npy"),
+                in_memory=in_memory,
+            ),
+        ]
+        feature_store = gb.TorchBasedFeatureStore(feature_data)
+
+        expected_feature_store_str = str(
+            """TorchBasedFeatureStore{(<OnDiskFeatureDataDomain.NODE: 'node'>, 'paper', 'a'): TorchBasedFeature(feature=tensor([[1, 2, 4],
+                                                        [2, 5, 3]]),
+                                        metadata={},
+                      ), (<OnDiskFeatureDataDomain.EDGE: 'edge'>, 'paper:cites:paper', 'b'): TorchBasedFeature(feature=tensor([[[1, 2],
+                                                         [3, 4]],
+                                                
+                                                        [[2, 5],
+                                                         [3, 4]]]),
+                                        metadata={},
+                      )}"""
+        )
+        assert str(feature_store) == expected_feature_store_str, print(
+            feature_store
+        )
+
+        a = b = feature_data = None
+        feature_store = expected_feature_store_str = None