scikit-hep · Mar 19, 2025
diff --git a/‎docs/reference/toctree.txt
+5 b/‎docs/reference/toctree.txt
+5
diff --git a/‎src/awkward/_backends/dispatch.py
+22 b/‎src/awkward/_backends/dispatch.py
+22
diff --git a/‎src/awkward/_kernels.py
+5 b/‎src/awkward/_kernels.py
+5
diff --git a/‎src/awkward/_lookup.py
+5 b/‎src/awkward/_lookup.py
+5
diff --git a/‎src/awkward/_nplikes/__init__.py
+19 b/‎src/awkward/_nplikes/__init__.py
+19
@@ -73,6 +73,11 @@
     generated/ak.to_packed
     generated/ak.copy
 
+ .. toctree::
+    :caption: Materializing virtual arrays
+
+    generated/ak.materialize
+
 .. toctree::
     :caption: Validity checking
 
 
@@ -4,9 +4,11 @@
 
 from collections.abc import Collection
 
+import awkward as ak
 from awkward._backends.backend import Backend
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpy_like import NumpyLike, NumpyMetadata
+from awkward._nplikes.virtual import VirtualArray
 from awkward._typing import Callable, TypeAlias, TypeVar, cast
 from awkward._util import UNSET, Sentinel
 
@@ -70,6 +72,7 @@ def common_backend(backends: Collection[Backend]) -> Backend:
 
 
 def backend_of_obj(obj, default: D | Sentinel = UNSET) -> Backend | D:
+    # the backend of virtual arrays will be determined via the `find_virtual_backend` lookup
     cls = type(obj)
     try:
         lookup = _type_to_backend_lookup[cls]
@@ -129,3 +132,22 @@ def regularize_backend(backend: str | Backend) -> Backend:
         return _name_to_backend_cls[backend].instance()
     else:
         raise ValueError(f"No such backend {backend!r} exists.")
+
+
+@register_backend_lookup_factory
+def find_virtual_backend(obj: type):
+    """
+    Implements a lookup for finding the backends of virtual arrays.
+    This is necessary to avoid calling `isinstance` inside `backend_of_obj` which may cause slowdown.
+    """
+    if issubclass(obj, VirtualArray):
+
+        def finder(obj: VirtualArray):
+            if isinstance(obj.nplike, ak._nplikes.numpy.Numpy):
+                return _name_to_backend_cls["cpu"].instance()
+            elif isinstance(obj.nplike, ak._nplikes.cupy.Cupy):
+                return _name_to_backend_cls["cuda"].instance()
+            else:
+                raise TypeError("A virtual array can only have numpy or cupy backends")
+
+        return finder
@@ -12,6 +12,7 @@
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._nplikes.typetracer import try_touch_data
+from awkward._nplikes.virtual import materialize_if_virtual
 from awkward._typing import Protocol, TypeAlias
 
 KernelKeyType: TypeAlias = tuple  # Tuple[str, Unpack[Tuple[metadata.dtype, ...]]]
@@ -88,6 +89,8 @@ def _cast(cls, x, t):
     def __call__(self, *args) -> None:
         assert len(args) == len(self._impl.argtypes)
 
+        args = materialize_if_virtual(*args)
+
         return self._impl(
             *(self._cast(x, t) for x, t in zip(args, self._impl.argtypes))
         )
@@ -138,6 +141,8 @@ def _cast(self, x, type_):
     def __call__(self, *args) -> None:
         import awkward._connect.cuda as ak_cuda
 
+        args = materialize_if_virtual(*args)
+
         cupy = ak_cuda.import_cupy("Awkward Arrays with CUDA")
         maxlength = self.max_length(args)
         grid, blocks = self.calc_grid(maxlength), self.calc_blocks(maxlength)
 
@@ -30,6 +30,11 @@ def arrayptr(x):
 
 
 def tolookup(layout, positions):
+    if not layout.is_all_materialized:
+        raise TypeError(
+            "Only fully materialized arrays can be passed into lookups. Use ak.materialize to materialize the array before passing it to the kernel."
+        )
+
     if isinstance(layout, ak.contents.EmptyArray):
         return tolookup(layout.to_NumpyArray(np.dtype(np.float64)), positions)
 
 
@@ -6,6 +6,7 @@
 import awkward._nplikes.jax
 import awkward._nplikes.numpy
 import awkward._nplikes.typetracer
+import awkward._nplikes.virtual
 from awkward._nplikes.dispatch import nplike_of_obj
 from awkward._typing import TYPE_CHECKING
 
@@ -27,6 +28,24 @@ def to_nplike(
     if from_nplike is nplike:
         return array
 
+    # We can always convert virtual arrays to typetracers
+    # but can only convert virtual arrays to other backends with known data if they are intentionally materialized
+    # Only numpy and cupy nplikes are allowed for virtual arrays
+    if isinstance(array, awkward._nplikes.virtual.VirtualArray):
+        if not array.is_materialized and nplike.known_data:
+            raise TypeError(
+                "Cannot convert a VirtualArray to a different nplike with known data without materializing it first. Use ak.materialize on the array to do so."
+            )
+        else:
+            if nplike.supports_virtual_arrays:
+                array = array.materialize()
+            elif not nplike.known_data:
+                pass
+            else:
+                raise TypeError(
+                    f"Can only convert a VirtualArray to numpy, cupy or typetracer nplikes. Received {type(nplike)}"
+                )
+
     if nplike.known_data and not from_nplike.known_data:
         raise TypeError(
             "Converting from an nplike without known data to an nplike with known data is not supported"