title: resample and groupby on big dask array with xarray - using map_blocks?
body: <p>I have a custom workflow, that requires using <code>resample</code> to get to a higher temporal frequency, applying a <code>ufunc</code>, and <code>groupby</code> + <code>mean</code> to compute the final result.</p>
  <p>I would like to apply this to a big <code>xarray</code> dataset, which is backed by a chunked <code>dask</code> array. For computation, I'd like to use <code>dask.distributed</code>.</p>
  <p>However, when I apply this to the full dataset, the number of tasks skyrockets, overwhelming the client and most likely also the scheduler and workers if submitted.</p>
  <p>The <a href="https://xarray-test.readthedocs.io/en/latest/dask.html" rel="nofollow noreferrer">xarray docs</a> explain:</p>
  <blockquote>
  <p>Do your spatial and temporal indexing (e.g. .sel() or .isel()) early
  in the pipeline, especially before calling resample() or groupby().
  Grouping and rasampling triggers some computation on all the blocks,
  which in theory should commute with indexing, but this optimization
  hasn’t been implemented in dask yet.</p>
  </blockquote>
  <p>But I really need to apply this to the <em>full</em> temporal axis.</p>
  <p>So how to best implement this?</p>
  <p>My approach was to use <code>map_blocks</code>, to apply this function for each chunk individually as to keep the individual <code>xarray</code> sub-datasets small enough.</p>
  <p>This seems to work on a small scale, but when I use the full dataset, the workers run out of memory and quickly die.</p>
  <p>Looking at the dashboard, the function I'm applying to the array gets executed multiple times of the number of chunks I have. Shouldn't these two numbers line up?</p>
  <p><strong>So my questions are</strong>:</p>
  <ul>
  <li>Is this approach valid?</li>
  <li>How could I implement this workflow otherwise, besides manually implementing the <code>resample</code> and <code>groupby</code> part and putting it in a <code>ufunc</code>?</li>
  <li>Any ideas regarding the performance issues at scale (specifically the number of executions vs chunks)?</li>
  </ul>
  <p>Here's a small example that mimics the workflow and shows the number of executions vs chunks:</p>
  <pre class="lang-py prettyprint-override"><code>from time import sleep
  
  import dask
  from dask.distributed import Client, LocalCluster
  import numpy as np
  import pandas as pd
  import xarray as xr
  
  def ufunc(x):
      # computation
      sleep(2)
      return x
  
  def fun(x):
      # upsample to higher res
      x = x.resample(time=&quot;1h&quot;).asfreq().fillna(0)
      
      # apply function
      x = xr.apply_ufunc(ufunc, x, input_core_dims=[[&quot;time&quot;]], output_core_dims=[['time']], dask=&quot;parallelized&quot;)
      
      # average over dates
      x['time'] = x.time.dt.strftime(&quot;%Y-%m-%d&quot;)
      x = x.groupby(&quot;time&quot;).mean()
  
      return x
  
  def create_xrds(shape):
      ''' helper function to create dataset'''
      x,y,t = shape
  
      tv = pd.date_range(start=&quot;1970-01-01&quot;, periods=t)
      
      ds = xr.Dataset({
      &quot;band&quot;: xr.DataArray(
          dask.array.zeros(shape, dtype=&quot;int16&quot;),
          dims=['x', 'y', 'time'],
          coords={&quot;x&quot;: np.arange(0, x), &quot;y&quot;: np.arange(0, y), &quot;time&quot;: tv})
      })
      
      return ds
  
  
  # set up distributed
  
  cluster = LocalCluster(n_workers=2)
  client = Client(cluster)
  
  ds = create_xrds((500,500,500)).chunk({&quot;x&quot;: 100, &quot;y&quot;: 100, &quot;time&quot;: -1})
  
  # create template
  
  template = ds.copy()
  template['time'] = template.time.dt.strftime(&quot;%Y-%m-%d&quot;)
  
  # map fun to blocks
  ds_out = xr.map_blocks(fun, ds, template=template)
  
  # persist
  
  ds_out.persist()
  </code></pre>
  <p>Using the example above, this is how the dask array (25 chunks) looks like:</p>
  <p><a href="https://i.stack.imgur.com/NvqX1.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/NvqX1.png" alt="enter image description here" /></a></p>
  <p>But the function <code>fun</code> gets executed 125 times:</p>
  <p><a href="https://i.stack.imgur.com/Mk0dB.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/Mk0dB.png" alt="enter image description here" /></a></p>