[GraphBolt] Get feature keys of FeatureStore. (#6920)

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

python/dgl/graphbolt/feature_store.py

@@ -171,3 +171,14 @@ class FeatureStore:
             updated.
         """
         raise NotImplementedError
+
+    def keys(self):
+        """Get the keys of the features.
+
+        Returns
+        -------
+        List[tuple]
+            The keys of the features. The tuples are in `(domain, type_name,
+            feat_name)` format.
+        """
+        raise NotImplementedError

python/dgl/graphbolt/impl/basic_feature_store.py

@@ -136,3 +136,14 @@ class BasicFeatureStore(FeatureStore):
     def __len__(self):
         """Return the number of features."""
         return len(self._features)
+
+    def keys(self):
+        """Get the keys of the features.
+
+        Returns
+        -------
+        List[tuple]
+            The keys of the features. The tuples are in `(domain, type_name,
+            feat_name)` format.
+        """
+        return list(self._features.keys())

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

@@ -43,6 +43,9 @@ def test_basic_feature_store_homo():
     assert feature_store.metadata("node", None, "a") == metadata
     assert feature_store.metadata("node", None, "b") == {}
 
+    # Test get keys of the features.
+    assert feature_store.keys() == [("node", None, "a"), ("node", None, "b")]
+
 
 def test_basic_feature_store_hetero():
     a = torch.tensor([[1, 2, 4], [2, 5, 3]])
@@ -81,6 +84,12 @@ def test_basic_feature_store_hetero():
     assert feature_store.metadata("node", "author", "a") == metadata
     assert feature_store.metadata("edge", "paper:cites", "b") == {}
 
+    # Test get keys of the features.
+    assert feature_store.keys() == [
+        ("node", "author", "a"),
+        ("edge", "paper:cites", "b"),
+    ]
+
 
 def test_basic_feature_store_errors():
     a = torch.tensor([3, 2, 1])

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

@@ -239,6 +239,12 @@ def test_torch_based_feature_store(in_memory):
             "edge", "paper:cites:paper", "b"
         ) == torch.Size([2, 2])
 
+        # Test get the keys of the features.
+        assert feature_store.keys() == [
+            ("node", "paper", "a"),
+            ("edge", "paper:cites:paper", "b"),
+        ]
+
         # For windows, the file is locked by the numpy.load. We need to delete
         # it before closing the temporary directory.
         a = b = None