Query data when using dask-awkward

[JuanBSLeite]

Dear experts,

I'm a starter on uproot.dask. I'm trying to apply some selections on data using it. Is there a way to apply selections to Dask-Awkward similar to Dask.dataframe.query?

Something like ddf.query('y > x').compute() ?

https://docs.dask.org/en/stable/generated/dask.dataframe.DataFrame.query.html

Thanks,
Juan

[jpivarski]

Yes, here's an example (given the array from here):

>>> import awkward as ak
>>> import dask_awkward as dak
>>> array = ak.Array([
...     [{"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}, {"x": 3.3, "y": [1, 2, 3]}],
...     [],
...     [{"x": 4.4, "y": [1, 2, 3, 4]}, {"x": 5.5, "y": [1, 2, 3, 4, 5]}]
... ])

>>> delayed_array = dak.from_awkward(array, npartitions=1)
>>> selected_delayed_array = delayed_array[delayed_array.x > 3]

>>> selected_delayed_array.compute().show()
[[{x: 3.3, y: [1, 2, 3]}],
 [],
 [{x: 4.4, y: [1, 2, 3, 4]}, {x: 5.5, y: [1, ..., 5]}]]

Only the records in which x > 3 remain, though the records with fields x, y are in variable-length lists and x and y have very different data types (scalar vs list).

[JuanBSLeite]

Thanks @jpivarski for the quick answer. But, I need to apply multiple selections using dataframe variables. Can I do selections like bellow?

selected_delayed_array = delayed_array[delayed_array.x > 3 and delayed_array.y>1]

[JuanBSLeite]

I found it where in the documentation. Thanks again!

[jpivarski]

Each delayed_array.x > 3 and delayed_array.y > 1 creates an array of booleans that can be used as a slice. To combine the two arrays of booleans with "and", use the bitwise operator & (with parentheses for correct order-of-operations):

selected_delayed_array = delayed_array[(delayed_array.x > 3) & (delayed_array.y > 1)]

[jpivarski]

In this particular case, x field values are scalars and y field values are lists, which might broadcast when used together in the & operation, but I'd have to think a bit more.

Your case might not have this feature: if you have as many x values as you have y values, there would be no issue.