<jupyter_start><jupyter_text>dask-awkward Proof-of-Concept: October 2021**_This notebook is using a very alpha version of dask-awkward; this is not a stable notebook! The features shown are also not fully tested._**<jupyter_code>import awkward as ak
import dask_awkward as dak
import dask_awkward.data<jupyter_output><empty_output><jupyter_text>I/OSo far we've implemented two very simple readers: JSON and Parquet. The JSON reader partitions the data such that each file or string in a list of sources will be a chunk/partition of the complete `DaskAwkwardArray` collection. The parquet reader can partition via a user provided `row_groups` argument or by using a list of files. For this small demo we'll use the JSON reader. We've created a simple `dask_awkward.data` module to provide some simple JSON data for testing/demoing:<jupyter_code>data = dask_awkward.data.json_data("numbers")
len(data)  # 3 JSON formatted strings
# This function is designed to be the dask version of ak.from_json
daa = dak.from_json(data)<jupyter_output><empty_output><jupyter_text>The `daa` variable will be our Dask awkward array.<jupyter_code># provides a DaskAwkwardArray collection
daa
daa.visualize()
# Sanity check; let's compare to the concrete awkward array version
a0 = ak.from_json(data[0])
a1 = ak.from_json(data[1])
a2 = ak.from_json(data[2])
ak.all(ak.concatenate([a0, a1, a2]) == daa.compute())<jupyter_output><empty_output><jupyter_text>The `caa` variable will be our concrete awkward array<jupyter_code>caa = ak.concatenate([a0, a1, a2])
ak.min(caa)
dak.min(daa)<jupyter_output><empty_output><jupyter_text>We currently have very minimal metadata support; in the case where `axis=None` is passed to `dak.min`, we know the result will be a scalar. We can see the constructed task graph:<jupyter_code>dak.min(daa).visualize()
dak.min(daa).compute()<jupyter_output><empty_output><jupyter_text>Now with an axis argument (fewer steps in the graph!):<jupyter_code>dak.min(daa, axis=1).visualize()
dak.min(daa, axis=1).compute()
ak.all(ak.min(caa, axis=1) == dak.min(daa, axis=1).compute())<jupyter_output><empty_output><jupyter_text>Now we'll look at some data with labels; the data is structured such that each element has an `analysis` record; where inside we have `x{1,2}`, `y{1,2}`, `z{1,2}`, `t{1,2}` records with random different lengths. We'll use this to showcase the attribute access functionality.<jupyter_code>daa = dak.from_json(dask_awkward.data.json_data(kind="records"))
daa
daa.compute().to_list()<jupyter_output><empty_output><jupyter_text>Let's define a function that squares the input:<jupyter_code>def sq(x):
    return x * x<jupyter_output><empty_output><jupyter_text>Now square all `x1`'s and all `y1`'s via the collection's `map_partitions` method.<jupyter_code>x1sq = daa.analysis.x1.map_partitions(sq)
y1sq = daa.analysis.y1.map_partitions(sq)<jupyter_output><empty_output><jupyter_text>Now the module level & more general `map_partitions` function can be used with multiple collections, let's map `np.add` to add all of the squares, then get the max of the result.<jupyter_code>import numpy as np

calc = dak.map_partitions(np.add, x1sq, y1sq)
res = dak.max(calc)
res.visualize()<jupyter_output><empty_output><jupyter_text>From the task graph visualization we can see all of the steps:- reading the JSON- `getitem` for the `.` access (`.analysis.x1` and `.analysis.y1`)- Our `sq` function- `add` from the last `map_partitions` call- finally, the `max` call - yielding multiple nodes in the graph (calculating max on the partitions followed up with a naïve reduction).<jupyter_code>res.compute()<jupyter_output><empty_output><jupyter_text>Sanity check (using the simplest form provided by the complete awkward API which we will eventually have in dask-awkward!):<jupyter_code>caa = daa.compute()
ak.max(caa.analysis.x1 ** 2 + caa.analysis.y1 ** 2)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup