Repository: pyg-team/pytorch_geometric
Pull request #7493 (2023-06-02 11:12:48)

`cudf` support in `map_index`


Patch:
diff --git a/.github/workflows/building_rusty1s_conda.yml b/.github/workflows/building_rusty1s_conda.yml
index 296a402cb..d0b88eda2 100644
--- a/.github/workflows/building_rusty1s_conda.yml
+++ b/.github/workflows/building_rusty1s_conda.yml
@@ -11,8 +11,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-18.04, macos-10.15, windows-2019]
-        # We have troube building for Python 3.11 due to version conflicts.
-        python-version: ['3.7', '3.8', '3.9', '3.10']  # '3.11'
+        python-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
         torch-version: [1.12.0, 1.13.0, 2.0.0]
         cuda-version: ['cpu', 'cu102', 'cu113', 'cu116', 'cu117', 'cu118']
         exclude:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f2c1e3f16..c1aecb3e9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- added a CPU-based and GPU-based `map_index` implementation ([#7493](https://github.com/pyg-team/pytorch_geometric/pull/7493))
 - Added the `AmazonBook` heterogeneous dataset ([#7483](https://github.com/pyg-team/pytorch_geometric/pull/7483))
 - Added hierarichial heterogeneous GraphSAGE example on OGB-MAG ([#7425](https://github.com/pyg-team/pytorch_geometric/pull/7425))
 - Added the `torch_geometric.distributed` package ([#7451](https://github.com/pyg-team/pytorch_geometric/pull/7451), [#7452](https://github.com/pyg-team/pytorch_geometric/pull/7452))
diff --git a/torch_geometric/utils/map.py b/torch_geometric/utils/map.py
index cea56b0ed..39db6520d 100644
--- a/torch_geometric/utils/map.py
+++ b/torch_geometric/utils/map.py
@@ -1,16 +1,29 @@
-from typing import Tuple
+import warnings
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
+from torch.utils.dlpack import from_dlpack
 
 
-def map_index(src: Tensor, index: Tensor) -> Tuple[Tensor, Tensor]:
+def map_index(
+    src: Tensor,
+    index: Tensor,
+    max_index: Optional[int] = None,
+    inclusive: bool = False,
+) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Maps indices in :obj:`src` to the positional value of their
     corresponding occurence in :obj:`index`.
+    Indices must be strictly positive.
 
     Args:
         src (torch.Tensor): The source tensor to map.
         index (torch.Tensor): The index tensor that denotes the new mapping.
+        max_index (int, optional): The maximum index value.
+            (default :obj:`None`)
+        inclusive (bool, optional): If set to :obj:`True`, it is assumed that
+            every entry in :obj:`src` has a valid entry in :obj:`index`.
+            Can speed-up computation. (default: :obj:`False`)
 
     :rtype: (:class:`torch.Tensor`, :class:`torch.BoolTensor`)
 
@@ -28,23 +41,106 @@ def map_index(src: Tensor, index: Tensor) -> Tuple[Tensor, Tensor]:
         >>> map_index(src, index)
         (tensor([1, 2, 2, 0]), tensor([True, True, False, True, True]))
     """
-    import pandas as pd
-
-    assert src.dim() == 1 and index.dim() == 1
-    assert not src.is_floating_point()
-    assert not index.is_floating_point()
-
-    arange = pd.RangeIndex(0, index.size(0))
-    df = pd.DataFrame(index=index.detach().cpu().numpy(), data={'out': arange})
-    ser = pd.Series(src.detach().cpu(), name='key')
-    result = df.merge(ser, how='right', left_index=True, right_on='key')
-    out = torch.from_numpy(result['out'].values).to(index.device)
-
-    if out.is_floating_point():
-        mask = torch.isnan(out).logical_not_()
-        out = out[mask].to(index.dtype)
-        return out, mask
-
-    out = out.to(index.dtype)
-    mask = torch.ones_like(out, dtype=torch.bool)
-    return out, mask
+    if src.is_floating_point():
+        raise ValueError(f"Expected 'src' to be an index (got '{src.dtype}')")
+    if index.is_floating_point():
+        raise ValueError(f"Expected 'index' to be an index (got "
+                         f"'{index.dtype}')")
+    if src.device != index.device:
+        raise ValueError(f"Both 'src' and 'index' must be on the same device "
+                         f"(got '{src.device}' and '{index.device}')")
+
+    if max_index is None:
+        max_index = max(src.max(), index.max())
+
+    # If the `max_index` is in a reasonable range, we can accelerate this
+    # operation by creating a helper vector to perform the mapping.
+    # NOTE This will potentially consumes a large chunk of memory
+    # (max_index=10 million => ~75MB), so we cap it at a reasonable size:
+    THRESHOLD = 40_000_000 if src.is_cuda else 10_000_000
+    if max_index <= THRESHOLD:
+        if inclusive:
+            assoc = src.new_empty(max_index + 1)
+        else:
+            assoc = src.new_full((max_index + 1, ), -1)
+        assoc[index] = torch.arange(index.numel(), dtype=src.dtype,
+                                    device=src.device)
+        out = assoc[src]
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], mask
+
+    WITH_CUDF = False
+    if src.is_cuda:
+        try:
+            import cudf
+            WITH_CUDF = True
+        except ImportError:
+            import pandas as pd
+            warnings.warn("Using CPU-based processing within 'map_index' "
+                          "which may cause slowdowns and device "
+                          "synchronization. Consider installing 'cudf' to "
+                          "accelerate computation")
+    else:
+        import pandas as pd
+
+    if not WITH_CUDF:
+        left_ser = pd.Series(src.cpu().numpy(), name='left_ser')
+        right_ser = pd.Series(
+            index=index.cpu().numpy(),
+            data=pd.RangeIndex(0, index.size(0)),
+            name='right_ser',
+        )
+
+        result = pd.merge(left_ser, right_ser, how='left', left_on='left_ser',
+                          right_index=True)
+
+        out = torch.from_numpy(result['right_ser'].values).to(index.device)
+
+        if out.is_floating_point() and inclusive:
+            raise ValueError("Found invalid entries in 'src' that do not have "
+                             "a corresponding entry in 'index'. Set "
+                             "`inclusive=False` to ignore these entries.")
+
+        if out.is_floating_point():
+            mask = torch.isnan(out).logical_not_()
+            out = out[mask].to(index.dtype)
+            return out, mask
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], mask
+
+    else:
+        left_ser = cudf.Series(src, name='left_ser')
+        right_ser = cudf.Series(
+            index=index,
+            data=cudf.RangeIndex(0, index.size(0)),
+            name='right_ser',
+        )
+
+        result = cudf.merge(left_ser, right_ser, how='left',
+                            left_on='left_ser', right_index=True, sort=True)
+
+        if inclusive:
+            try:
+                out = from_dlpack(result['right_ser'].to_dlpack())
+            except ValueError:
+                raise ValueError("Found invalid entries in 'src' that do not "
+                                 "have a corresponding entry in 'index'. Set "
+                                 "`inclusive=False` to ignore these entries.")
+        else:
+            out = from_dlpack(result['right_ser'].fillna(-1).to_dlpack())
+
+        out = out[src.argsort().argsort()]  # Restore original order.
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], mask


Test patch:
diff --git a/test/utils/test_map.py b/test/utils/test_map.py
index 923bb5b81..df6eb6e54 100644
--- a/test/utils/test_map.py
+++ b/test/utils/test_map.py
@@ -1,21 +1,77 @@
+import pytest
 import torch
 
+from torch_geometric.profile import benchmark
+from torch_geometric.testing import withCUDA
 from torch_geometric.utils.map import map_index
 
 
-def test_map_index():
-    src = torch.tensor([2, 0, 1, 0, 3])
-    index = torch.tensor([3, 2, 0, 1])
+@withCUDA
+@pytest.mark.parametrize('max_index', [3, 100_000_000])
+def test_map_index(device, max_index):
+    src = torch.tensor([2, 0, 1, 0, max_index], device=device)
+    index = torch.tensor([max_index, 2, 0, 1], device=device)
 
-    out, mask = map_index(src, index)
+    out, mask = map_index(src, index, inclusive=True)
+    assert out.device == device
+    assert mask is None
     assert out.tolist() == [1, 2, 3, 2, 0]
-    assert mask.tolist() == [True, True, True, True, True]
 
 
-def test_map_index_na():
-    src = torch.tensor([2, 0, 1, 0, 3])
-    index = torch.tensor([3, 2, 0])
+@withCUDA
+@pytest.mark.parametrize('max_index', [3, 100_000_000])
+def test_map_index_na(device, max_index):
+    src = torch.tensor([2, 0, 1, 0, max_index], device=device)
+    index = torch.tensor([max_index, 2, 0], device=device)
 
-    out, mask = map_index(src, index)
+    out, mask = map_index(src, index, inclusive=False)
+    assert out.device == device
+    assert mask.device == device
     assert out.tolist() == [1, 2, 2, 0]
     assert mask.tolist() == [True, True, False, True, True]
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cuda')
+    args = parser.parse_args()
+
+    src = torch.randint(0, 100_000_000, (100_000, ), device=args.device)
+    index = src.unique()
+
+    def trivial_map(src, index, max_index, inclusive):
+        if max_index is None:
+            max_index = max(src.max(), index.max())
+
+        if inclusive:
+            assoc = src.new_empty(max_index + 1)
+        else:
+            assoc = src.new_full((max_index + 1, ), -1)
+        assoc[index] = torch.arange(index.numel(), device=index.device)
+        out = assoc[src]
+
+        if inclusive:
+            return out, None
+        else:
+            mask = out != -1
+            return out[mask], mask
+
+    print('Inclusive:')
+    benchmark(
+        funcs=[trivial_map, map_index],
+        func_names=['trivial', 'map_index'],
+        args=(src, index, None, True),
+        num_steps=100,
+        num_warmups=50,
+    )
+
+    print('Exclusive:')
+    benchmark(
+        funcs=[trivial_map, map_index],
+        func_names=['trivial', 'map_index'],
+        args=(src, index[:50_000], None, False),
+        num_steps=100,
+        num_warmups=50,
+    )