[GraphBolt] Feature fetcher base datapipe (#5908)

Co-authored-by: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>

python/dgl/graphbolt/__init__.py

@@ -9,6 +9,7 @@ from .graph_storage import *
 from .itemset import *
 from .minibatch_sampler import *
 from .feature_store import *
+from .feature_fetcher import *
 from .copy_to import *
 from .dataset import *
 from .subgraph_sampler import *

python/dgl/graphbolt/feature_fetcher.py

+"""Feature fetchers"""
+
+from torchdata.datapipes.iter import Mapper
+
+
+class FeatureFetcher(Mapper):
+    """Base feature fetcher.
+
+    This is equivalent to the following iterator:
+
+    .. code:: python
+
+       for data in datapipe:
+           yield feature_fetch_func(data)
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    fn : callable
+        The function that performs feature fetching.
+    """

tests/python/pytorch/graphbolt/test_feature_fetcher.py

+import dgl
+import dgl.graphbolt
+import pytest
+import torch
+
+
+def get_graphbolt_fetch_func():
+    feature_store = dgl.graphbolt.feature_store.TorchBasedFeatureStore(
+        {
+            "feature": torch.randn(200, 4),
+            "label": torch.randint(0, 10, (200,)),
+        }
+    )
+
+    def fetch_func(data):
+        return feature_store.read("feature", data), feature_store.read(
+            "label", data
+        )
+
+    return fetch_func
+
+
+def get_tensor_fetch_func():
+    feature_store = torch.randn(200, 4)
+    label = torch.randint(0, 10, (200,))
+
+    def fetch_func(data):
+        return feature_store[data], label[data]
+
+    return fetch_func
+
+
+@pytest.mark.parametrize(
+    "fetch_func", [get_graphbolt_fetch_func(), get_tensor_fetch_func()]
+)
+def test_FeatureFetcher(fetch_func):
+    itemset = dgl.graphbolt.ItemSet(torch.arange(10))
+    minibatch_dp = dgl.graphbolt.MinibatchSampler(itemset, batch_size=2)
+    fetcher_dp = dgl.graphbolt.FeatureFetcher(minibatch_dp, fetch_func)
+
+    assert len(list(fetcher_dp)) == 5