CuPy buffers error triggered from dask_awkward.from_parquet

[NJManganelli]

A file with a perhaps ususual (or bad) parquet schema is triggering an awkward error when computed on via dask_awkward.from_parquet.

# all imports included in the event it takes more than just dak to reproduce
import coffea
import numpy as np
import numba as nb
import awkward as ak
import dask_awkward as dak
import dask
import uproot
import hist
from hist.dask import Hist
import trino
import matplotlib.pyplot as plt
import s3fs
import fsspec
from minio import Minio
from coffea.nanoevents import NanoEventsFactory, BaseSchema


joined = dak.from_parquet("join_out_cupy_error_triggered.parquet")
ak.num(joined.electrons, axis=1).compute()

results in the following
(coffea: 2024.2.1 ak: 2.6.1 uproot: 5.2.2 hist: 2.7.2 dask-awkward: 2024.2.0 numpy: 1.26.4 numba: 0.59.0 trino: 0.327.0 as package versions)

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[32], line 2
      1 joined3 = dak.from_parquet("join_out_cupy_error_triggered.parquet")
----> 2 ak.num(joined3.electrons, axis=1).compute()

File [/usr/local/lib/python3.11/site-packages/dask/base.py:379](http://localhost:8190/usr/local/lib/python3.11/site-packages/dask/base.py#line=378), in DaskMethodsMixin.compute(self, **kwargs)
    355 def compute(self, **kwargs):
    356     """Compute this dask collection
    357 
    358     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    377     dask.compute
    378     """
--> 379     (result,) = compute(self, traverse=False, **kwargs)
    380     return result

File [/usr/local/lib/python3.11/site-packages/dask/base.py:665](http://localhost:8190/usr/local/lib/python3.11/site-packages/dask/base.py#line=664), in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    662     postcomputes.append(x.__dask_postcompute__())
    664 with shorten_traceback():
--> 665     results = schedule(dsk, keys, **kwargs)
    667 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File [/usr/local/lib/python3.11/site-packages/dask_awkward/lib/core.py:1890](http://localhost:8190/usr/local/lib/python3.11/site-packages/dask_awkward/lib/core.py#line=1889), in ArgsKwargsPackedFunction.__call__(self, *args_deps_expanded)
   1888     len_args += n_args
   1889 kwargs = self.kwarg_repacker(args_deps_expanded[len_args:])[0]
-> 1890 return self.fn(*args, **kwargs)

File [/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py:62](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_dispatch.py#line=61), in named_high_level_function.<locals>.dispatch(*args, **kwargs)
     60 # Failed to find a custom overload, so resume the original function
     61 try:
---> 62     next(gen_or_result)
     63 except StopIteration as err:
     64     return err.value

File [/usr/local/lib/python3.11/site-packages/awkward/operations/ak_num.py:83](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/operations/ak_num.py#line=82), in num(array, axis, highlevel, behavior, attrs)
     80 yield (array,)
     82 # Implementation
---> 83 return _impl(array, axis, highlevel, behavior, attrs)

File [/usr/local/lib/python3.11/site-packages/awkward/operations/ak_num.py:110](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/operations/ak_num.py#line=109), in _impl(array, axis, highlevel, behavior, attrs)
    107     elif layout.is_leaf:
    108         raise AxisError(f"axis={axis} exceeds the depth of this array ({depth})")
--> 110 out = ak._do.recursively_apply(layout, action, numpy_to_regular=True)
    112 return ctx.wrap(out, highlevel=highlevel)

File [/usr/local/lib/python3.11/site-packages/awkward/_do.py:36](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_do.py#line=35), in recursively_apply(layout, action, behavior, depth_context, lateral_context, allow_records, keep_parameters, numpy_to_regular, return_simplified, return_array, function_name, regular_to_jagged)
     21 def recursively_apply(
     22     layout: Content | Record,
     23     action: ActionType,
   (...)
     33     regular_to_jagged=False,
     34 ) -> Content | Record | None:
     35     if isinstance(layout, Content):
---> 36         return layout._recursively_apply(
     37             action,
     38             1,
     39             copy.copy(depth_context),
     40             lateral_context,
     41             {
     42                 "allow_records": allow_records,
     43                 "keep_parameters": keep_parameters,
     44                 "numpy_to_regular": numpy_to_regular,
     45                 "regular_to_jagged": regular_to_jagged,
     46                 "return_simplified": return_simplified,
     47                 "return_array": return_array,
     48                 "function_name": function_name,
     49             },
     50         )
     52     elif isinstance(layout, Record):
     53         out = recursively_apply(
     54             layout._array,
     55             action,
   (...)
     64             function_name,
     65         )

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py:770](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py#line=769), in BitMaskedArray._recursively_apply(self, action, depth, depth_context, lateral_context, options)
    768     return result
    769 elif result is None:
--> 770     return continuation()
    771 else:
    772     raise AssertionError(result)

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py:733](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py#line=732), in BitMaskedArray._recursively_apply.<locals>.continuation()
    730 def continuation():
    731     return make(
    732         self._mask,
--> 733         content._recursively_apply(
    734             action,
    735             depth,
    736             copy.copy(depth_context),
    737             lateral_context,
    738             options,
    739         ),
    740         self._valid_when,
    741         self._length,
    742         self._lsb_order,
    743         parameters=self._parameters if options["keep_parameters"] else None,
    744     )

File [/usr/local/lib/python3.11/site-packages/awkward/contents/listoffsetarray.py:2169](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/listoffsetarray.py#line=2168), in ListOffsetArray._recursively_apply(self, action, depth, depth_context, lateral_context, options)
   2165     offsetsmin = self._offsets[0]
   2166     offsets = ak.index.Index(
   2167         self._offsets.data - offsetsmin, nplike=self._backend.index_nplike
   2168     )
-> 2169     content = self._content[offsetsmin : self._offsets[-1]]
   2170 else:
   2171     self._touch_data(recursive=False)

File [/usr/local/lib/python3.11/site-packages/awkward/contents/content.py:519](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/content.py#line=518), in Content.__getitem__(self, where)
    518 def __getitem__(self, where):
--> 519     return self._getitem(where)

File [/usr/local/lib/python3.11/site-packages/awkward/contents/content.py:530](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/content.py#line=529), in Content._getitem(self, where)
    525 elif isinstance(where, slice) and where.step is None:
    526     # Ensure that start, stop are non-negative!
    527     start, stop, _, _ = self._backend.index_nplike.derive_slice_for_length(
    528         normalize_slice(where, nplike=self._backend.index_nplike), self.length
    529     )
--> 530     return self._getitem_range(start, stop)
    532 elif isinstance(where, slice):
    533     return self._getitem((where,))

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py:498](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py#line=497), in BitMaskedArray._getitem_range(self, start, stop)
    497 def _getitem_range(self, start: IndexType, stop: IndexType) -> Content:
--> 498     return self.to_ByteMaskedArray()._getitem_range(start, stop)

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bytemaskedarray.py:400](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bytemaskedarray.py#line=399), in ByteMaskedArray._getitem_range(self, start, stop)
    395     self._touch_shape(recursive=False)
    396     return self
    398 return ByteMaskedArray(
    399     self._mask[start:stop],
--> 400     self._content._getitem_range(start, stop),
    401     self._valid_when,
    402     parameters=self._parameters,
    403 )

File [/usr/local/lib/python3.11/site-packages/awkward/contents/recordarray.py:466](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/recordarray.py#line=465), in RecordArray._getitem_range(self, start, stop)
    457     return RecordArray(
    458         [],
    459         self._fields,
   (...)
    462         backend=self._backend,
    463     )
    464 else:
    465     return RecordArray(
--> 466         [x._getitem_range(start, stop) for x in self._contents],
    467         self._fields,
    468         length,
    469         parameters=self._parameters,
    470         backend=self._backend,
    471     )

File [/usr/local/lib/python3.11/site-packages/awkward/contents/recordarray.py:466](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/recordarray.py#line=465), in <listcomp>(.0)
    457     return RecordArray(
    458         [],
    459         self._fields,
   (...)
    462         backend=self._backend,
    463     )
    464 else:
    465     return RecordArray(
--> 466         [x._getitem_range(start, stop) for x in self._contents],
    467         self._fields,
    468         length,
    469         parameters=self._parameters,
    470         backend=self._backend,
    471     )

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py:498](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py#line=497), in BitMaskedArray._getitem_range(self, start, stop)
    497 def _getitem_range(self, start: IndexType, stop: IndexType) -> Content:
--> 498     return self.to_ByteMaskedArray()._getitem_range(start, stop)

File [/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py:396](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/contents/bitmaskedarray.py#line=395), in BitMaskedArray.to_ByteMaskedArray(self)
    388 bytemask = ak.index.Index8.empty(
    389     self._mask.length * 8, self._backend.index_nplike
    390 )
    391 assert (
    392     bytemask.nplike is self._backend.nplike
    393     and self._mask.nplike is self._backend.nplike
    394 )
    395 self._backend.maybe_kernel_error(
--> 396     self._backend[
    397         "awkward_BitMaskedArray_to_ByteMaskedArray",
    398         bytemask.dtype.type,
    399         self._mask.dtype.type,
    400     ](
    401         bytemask.data,
    402         self._mask.data,
    403         self._mask.length,
    404         False,  # this differs from the kernel call in 'bytemask'
    405         self._lsb_order,
    406     )
    407 )
    408 return ByteMaskedArray(
    409     bytemask[: self._length],
    410     self._content,
    411     self._valid_when,
    412     parameters=self._parameters,
    413 )

File [/usr/local/lib/python3.11/site-packages/awkward/_kernels.py:93](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_kernels.py#line=92), in NumpyKernel.__call__(self, *args)
     90 def __call__(self, *args) -> None:
     91     assert len(args) == len(self._impl.argtypes)
---> 93     return self._impl(
     94         *(self._cast(x, t) for x, t in zip(args, self._impl.argtypes))
     95     )

File [/usr/local/lib/python3.11/site-packages/awkward/_kernels.py:94](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_kernels.py#line=93), in <genexpr>(.0)
     90 def __call__(self, *args) -> None:
     91     assert len(args) == len(self._impl.argtypes)
     93     return self._impl(
---> 94         *(self._cast(x, t) for x, t in zip(args, self._impl.argtypes))
     95     )

File [/usr/local/lib/python3.11/site-packages/awkward/_kernels.py:84](http://localhost:8190/usr/local/lib/python3.11/site-packages/awkward/_kernels.py#line=83), in NumpyKernel._cast(cls, x, t)
     82         return ctypes.cast(x, t)
     83     else:
---> 84         raise AssertionError(
     85             "CuPy buffers shouldn't be passed to Numpy Kernels."
     86         )
     87 else:
     88     return x

AssertionError: CuPy buffers shouldn't be passed to Numpy Kernels.


See if this has been reported at https://github.com/scikit-hep/awkward/issues

When the equivalent is done via awkward.from_parquet, this computes without an error, but nevertheless it might belong properly there instead of dak.

The file can be found here:
https://drive.google.com/file/d/135ILkDokbHx4YehZuFSZ455hyRtBmfmE/view?usp=sharing

[agoose77]

Hmm, this is not a very helpful error message. It looks like we're treating

if x:
   ...
else:
   ...

as a boolean "NumPy" or "CuPy" rather than "not NumPy". I suspect we're seeing placeholders here, but they shouldn't be getting in to a kernel like this. Could you share the analysis excerpt too?

[NJManganelli]

The two lines at the bottom of the first code-block are all that's needed, I believe. The schema for this file, read via pyarrow.parquet.read_schema:

run: int32
event: int32
luminosityblock: int32
electrons: list<element: struct<pt: float, eta: float, phi: float, mass: float>>
  child 0, element: struct<pt: float, eta: float, phi: float, mass: float>
      child 0, pt: float
      child 1, eta: float
      child 2, phi: float
      child 3, mass: float
gnn_p1: float
gnn_p2: float
gnn_p3: float
gnn_p4: float
-- schema metadata --
writer.time.zone: 'Etc/GMT'

Edit for posterity: setting optimize_graph=False in the compute removes this triggered error message, just like the error with the other schema in the iris-hep:servicex-join channel being discussed now