feat: ability to optionally disable the caching of virtual arrays on themselves

src/awkward/_nplikes/array_module.py

@@ -80,6 +80,7 @@ def asarray(
                     obj._dtype if dtype is None else dtype,
                     lambda: self.asarray(obj.materialize(), dtype=dtype, copy=copy),
                     lambda: obj.shape,
+                    __enable_caching__=obj.__enable_caching__,
                 )
         if copy:
             return self._module.array(obj, dtype=dtype, copy=True)
@@ -108,6 +109,7 @@ def ascontiguousarray(
                     x._dtype,
                     lambda: self.ascontiguousarray(x.materialize()),  # type: ignore[arg-type]
                     lambda: x.shape,
+                    __enable_caching__=x.__enable_caching__,
                 )
         else:
             return self._module.ascontiguousarray(x)
@@ -356,6 +358,7 @@ def reshape(
                     x.dtype,
                     lambda: self.reshape(x.materialize(), next_shape, copy=copy),  # type: ignore[arg-type]
                     None,
+                    __enable_caching__=x.__enable_caching__,
                 )
 
         if copy is None:

src/awkward/_nplikes/virtual.py

@@ -81,6 +81,7 @@ def __init__(
         generator: Callable[[], ArrayLike],
         shape_generator: Callable[[], tuple[ShapeItem, ...]] | None = None,
         __wrap_generator_asarray__: bool = False,
+        __enable_caching__: bool = True,
     ) -> None:
         if not nplike.supports_virtual_arrays:
             raise TypeError(
@@ -104,6 +105,8 @@ def __init__(
         self._generator = generator
         self._shape_generator = shape_generator
 
+        self.__enable_caching__ = __enable_caching__
+
     @property
     def dtype(self) -> DType:
         return self._dtype
@@ -180,9 +183,11 @@ def materialize(self) -> ArrayLike:
                     f"{type(self).__name__} had dtype {self._dtype} before materialization while the materialized array has dtype {array.dtype}"
                 )
             self._shape = array.shape
-            self._array = array
-            self._shape_generator = assert_never
-            self._generator = assert_never
+            if self.__enable_caching__:
+                self._array = array
+                self._shape_generator = assert_never
+                self._generator = assert_never
+            return array
         return self._array  # type: ignore[return-value]
 
     @property
@@ -205,6 +210,7 @@ def T(self):
             self._dtype,
             lambda: self.materialize().T,
             lambda: self.shape[::-1],
+            __enable_caching__=self.__enable_caching__,
         )
 
     def view(self, dtype: DTypeLike) -> Self:
@@ -237,6 +243,7 @@ def view(self, dtype: DTypeLike) -> Self:
             dtype,
             lambda: self.materialize().view(dtype),
             None,
+            __enable_caching__=self.__enable_caching__,
         )
 
     @property
@@ -263,6 +270,7 @@ def byteswap(self, inplace=False):
             self._dtype,
             lambda: self.materialize().byteswap(inplace=inplace),
             lambda: self.shape,
+            __enable_caching__=self.__enable_caching__,
         )
 
     def tobytes(self, order="C") -> bytes:
@@ -275,6 +283,7 @@ def __copy__(self) -> VirtualNDArray:
             self._dtype,
             self._generator,
             self._shape_generator,
+            __enable_caching__=self.__enable_caching__,
         )
         new_virtual._array = self._array
         return new_virtual
@@ -287,6 +296,7 @@ def __deepcopy__(self, memo) -> VirtualNDArray:
             self._dtype,
             lambda: copy.deepcopy(current_generator(), memo),
             self._shape_generator,
+            __enable_caching__=self.__enable_caching__,
         )
         new_virtual._array = (
             copy.deepcopy(self._array, memo)
@@ -340,6 +350,7 @@ def __getitem__(self, index):
                 self._dtype,
                 lambda: self.materialize()[index],
                 None,
+                __enable_caching__=self.__enable_caching__,
             )
         else:
             return self.materialize().__getitem__(index)

src/awkward/_pickle.py

@@ -115,6 +115,7 @@ def unpickle_array_schema_1(
         buffer_key="{form_key}-{attribute}",
         byteorder="<",
         simplify=False,
+        enable_virtualarray_caching=True,
     )
 
 
@@ -141,6 +142,7 @@ def unpickle_record_schema_1(
         buffer_key="{form_key}-{attribute}",
         byteorder="<",
         simplify=False,
+        enable_virtualarray_caching=True,
     )
     layout = LowLevelRecord(array_layout, at)
     return Record(layout, behavior=behavior, attrs=attrs)

src/awkward/forms/form.py

@@ -540,6 +540,7 @@ def length_zero_array(
             behavior=behavior,
             attrs=None,
             simplify=False,
+            enable_virtualarray_caching=True,
         )
 
     def length_one_array(
@@ -696,6 +697,7 @@ def prepare(form, multiplier):
             behavior=behavior,
             attrs=None,
             simplify=False,
+            enable_virtualarray_caching=True,
         )
 
     def _expected_from_buffers(

src/awkward/highlevel.py

@@ -2953,6 +2953,7 @@ def snapshot(self):
                 backend="cpu",
                 byteorder=ak._util.native_byteorder,
                 simplify=True,
+                enable_virtualarray_caching=True,
                 highlevel=True,
                 behavior=self._behavior,
                 attrs=self._attrs,

src/awkward/operations/ak_from_avro_file.py

@@ -73,5 +73,6 @@ def _impl(form, length, container, highlevel, behavior, attrs):
         highlevel=highlevel,
         behavior=behavior,
         simplify=True,
+        enable_virtualarray_caching=True,
         attrs=attrs,
     )