content: from datashader.compiler import compile_components
  from datashader.utils import Dispatcher
  from datashader.glyphs.line import LinesXarrayCommonX
  from datashader.glyphs.quadmesh import (
      QuadMeshRaster, QuadMeshRectilinear, QuadMeshCurvilinear, build_scale_translate
  )
  from datashader.utils import apply
  import dask
  import numpy as np
  import xarray as xr
  from dask.base import tokenize, compute
  from dask.array.overlap import overlap
  dask_glyph_dispatch = Dispatcher()
  
  
  def dask_xarray_pipeline(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
      dsk, name = dask_glyph_dispatch(
          glyph, xr_ds, schema, canvas, summary, antialias=antialias, cuda=cuda)
  
      # Get user configured scheduler (if any), or fall back to default
      # scheduler for dask DataFrame
      scheduler = dask.base.get_scheduler() or xr_ds.__dask_scheduler__
      keys = xr_ds.__dask_keys__()
      optimize = xr_ds.__dask_optimize__
      graph = xr_ds.__dask_graph__()
  
      dsk.update(optimize(graph, keys))
  
      return scheduler(dsk, name)
  
  
  def shape_bounds_st_and_axis(xr_ds, canvas, glyph):
      if not canvas.x_range or not canvas.y_range:
          x_extents, y_extents = glyph.compute_bounds_dask(xr_ds)
      else:
          x_extents, y_extents = None, None
  
      x_range = canvas.x_range or x_extents
      y_range = canvas.y_range or y_extents
      x_min, x_max, y_min, y_max = bounds = compute(*(x_range + y_range))
      x_range, y_range = (x_min, x_max), (y_min, y_max)
  
      width = canvas.plot_width
      height = canvas.plot_height
  
      x_st = canvas.x_axis.compute_scale_and_translate(x_range, width)
      y_st = canvas.y_axis.compute_scale_and_translate(y_range, height)
      st = x_st + y_st
      shape = (height, width)
  
      x_axis = canvas.x_axis.compute_index(x_st, width)
      y_axis = canvas.y_axis.compute_index(y_st, height)
      axis = dict([(glyph.x_label, x_axis), (glyph.y_label, y_axis)])
  
      return shape, bounds, st, axis
  
  
  def dask_rectilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
      shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
  
      # Compile functions
      create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
          compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
      x_mapper = canvas.x_axis.mapper
      y_mapper = canvas.y_axis.mapper
      extend = glyph._build_extend(
          x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
      x_range = bounds[:2]
      y_range = bounds[2:]
  
      # Build chunk indices for coordinates
      chunk_inds = {}
      for k, chunks in xr_ds.chunks.items():
          chunk_inds[k] = [0] + list(np.cumsum(chunks))
  
      x_name = glyph.x
      y_name = glyph.y
      coords = xr_ds[glyph.name].coords
      coord_dims = list(coords.dims)
      xdim_ind = coord_dims.index(x_name)
      ydim_ind = coord_dims.index(y_name)
  
      var_name = list(xr_ds.data_vars.keys())[0]
  
      # Compute interval breaks
      xs = xr_ds[x_name].values
      ys = xr_ds[y_name].values
      x_breaks = glyph.infer_interval_breaks(xs)
      y_breaks = glyph.infer_interval_breaks(ys)
  
      def chunk(np_arr, *inds):
          # Reconstruct dataset for chunk from numpy array and chunk indices
          chunk_coords_list = []
          # for i, (coord_name, coord_vals) in enumerate(coords.items()):
          for i, coord_name in enumerate(coords.dims):
              chunk_number = inds[i]
              coord_slice = slice(
                  chunk_inds[coord_name][chunk_number],
                  chunk_inds[coord_name][chunk_number + 1]
              )
              chunk_coords_list.append([coord_name, coords[coord_name][coord_slice]])
  
          chunk_coords = dict(chunk_coords_list)
          chunk_ds = xr.DataArray(
              np_arr, coords=chunk_coords, dims=coord_dims, name=var_name
          ).to_dataset()
  
          # Compute chunk x/y breaks
          x_chunk_number = inds[xdim_ind]
          x_breaks_slice = slice(
              chunk_inds[x_name][x_chunk_number],
              chunk_inds[x_name][x_chunk_number + 1] + 1
          )
          x_breaks_chunk = x_breaks[x_breaks_slice]
  
          y_chunk_number = inds[ydim_ind]
          y_breaks_slice = slice(
              chunk_inds[y_name][y_chunk_number],
              chunk_inds[y_name][y_chunk_number + 1] + 1
          )
          y_breaks_chunk = y_breaks[y_breaks_slice]
  
          # Initialize aggregation buffers
          aggs = create(shape)
  
          # Perform aggregation
          extend(aggs, chunk_ds, st, bounds,
                 x_breaks=x_breaks_chunk, y_breaks=y_breaks_chunk)
          return aggs
  
      name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
      keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
      keys2 = [(name, i) for i in range(len(keys))]
      dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
      dsk[name] = (apply, finalize, [(combine, keys2)],
                   dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
                        attrs=dict(x_range=x_range, y_range=y_range)))
      return dsk, name
  
  
  def dask_raster(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
      shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
  
      # Compile functions
      create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
          compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
      x_mapper = canvas.x_axis.mapper
      y_mapper = canvas.y_axis.mapper
      extend = glyph._build_extend(
          x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
      x_range = bounds[:2]
      y_range = bounds[2:]
  
      # Build chunk indices for coordinates
      chunk_inds = {}
      for k, chunks in xr_ds.chunks.items():
          chunk_inds[k] = [0] + list(np.cumsum(chunks))
  
      x_name = glyph.x
      y_name = glyph.y
  
      coords = xr_ds[glyph.name].coords
  
      coord_dims = list(coords.dims)
      xdim_ind = coord_dims.index(x_name)
      ydim_ind = coord_dims.index(y_name)
      var_name = list(xr_ds.data_vars.keys())[0]
  
      # Pre-compute bin sizes. We do this here to handle length-1 chunks
      src_x0, src_x1 = glyph. _compute_bounds_from_1d_centers(
          xr_ds, x_name, maybe_expand=False, orient=False
      )
      src_y0, src_y1 = glyph._compute_bounds_from_1d_centers(
          xr_ds, y_name, maybe_expand=False, orient=False
      )
      xbinsize = abs(float(xr_ds[x_name][1] - xr_ds[x_name][0]))
      ybinsize = abs(float(xr_ds[y_name][1] - xr_ds[y_name][0]))
  
      # Compute scale/translate
      out_h, out_w = shape
      src_h, src_w = [xr_ds[glyph.name].shape[i] for i in [ydim_ind, xdim_ind]]
      out_x0, out_x1, out_y0, out_y1 = bounds
      scale_y, translate_y = build_scale_translate(
          out_h, out_y0, out_y1, src_h, src_y0, src_y1
      )
  
      scale_x, translate_x = build_scale_translate(
          out_w, out_x0, out_x1, src_w, src_x0, src_x1
      )
  
      def chunk(np_arr, *inds):
          # Reconstruct dataset for chunk from numpy array and chunk indices
          chunk_coords_list = []
          for i, coord_name in enumerate(coords.dims):
              chunk_number = inds[i]
              coord_slice = slice(
                  chunk_inds[coord_name][chunk_number],
                  chunk_inds[coord_name][chunk_number + 1]
              )
              chunk_coords_list.append([coord_name, coords[coord_name][coord_slice]])
  
          chunk_coords = dict(chunk_coords_list)
          chunk_ds = xr.DataArray(
              np_arr, coords=chunk_coords, dims=coord_dims, name=var_name
          ).to_dataset()
  
          # Compute offsets
          x_chunk_number = inds[xdim_ind]
          offset_x = chunk_inds[x_name][x_chunk_number]
  
          y_chunk_number = inds[ydim_ind]
          offset_y = chunk_inds[y_name][y_chunk_number]
  
          # Initialize aggregation buffers
          aggs = create(shape)
  
          # Perform aggregation
          extend(aggs, chunk_ds, st, bounds,
                 scale_x=scale_x, scale_y=scale_y,
                 translate_x=translate_x, translate_y=translate_y,
                 offset_x=offset_x, offset_y=offset_y,
                 src_xbinsize=xbinsize, src_ybinsize=ybinsize)
  
          return aggs
  
      name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
      keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
      keys2 = [(name, i) for i in range(len(keys))]
      dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
      dsk[name] = (apply, finalize, [(combine, keys2)],
                   dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
                        attrs=dict(x_range=x_range, y_range=y_range)))
      return dsk, name
  
  
  def dask_curvilinear(glyph, xr_ds, schema, canvas, summary, *, antialias=False, cuda=False):
      shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
  
      # Compile functions
      create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, _ = \
          compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda, partitioned=True)
      x_mapper = canvas.x_axis.mapper
      y_mapper = canvas.y_axis.mapper
      extend = glyph._build_extend(
          x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
      x_range = bounds[:2]
      y_range = bounds[2:]
  
      x_coord_name = glyph.x
      y_coord_name = glyph.y
      z_name = glyph.name
  
      data_dim_names = list(xr_ds[z_name].dims)
      x_coord_dim_names = list(xr_ds[x_coord_name].dims)
      y_coord_dim_names = list(xr_ds[y_coord_name].dims)
      zs = xr_ds[z_name].data
      x_centers = xr_ds[glyph.x].data
      y_centers = xr_ds[glyph.y].data
  
      var_name = list(xr_ds.data_vars.keys())[0]
  
      # Validate coordinates
      err_msg = (
          "DataArray {name} is backed by a Dask array, \n"
          "but coordinate {coord} is not backed by a Dask array with identical \n"
          "dimension order and chunks"
      )
      if (not isinstance(x_centers, dask.array.Array) or
              xr_ds[glyph.name].dims != xr_ds[glyph.x].dims or
              xr_ds[glyph.name].chunks != xr_ds[glyph.x].chunks):
          raise ValueError(err_msg.format(name=glyph.name, coord=glyph.x))
  
      if (not isinstance(y_centers, dask.array.Array) or
              xr_ds[glyph.name].dims != xr_ds[glyph.y].dims or
              xr_ds[glyph.name].chunks != xr_ds[glyph.y].chunks):
          raise ValueError(err_msg.format(name=glyph.name, coord=glyph.y))
  
      # Make sure coordinates are floats so that overlap with nan will behave properly
      if x_centers.dtype.kind != 'f':
          x_centers = x_centers.astype(np.float64)
      if y_centers.dtype.kind != 'f':
          y_centers = y_centers.astype(np.float64)
  
      x_overlapped_centers = overlap(x_centers, depth=1, boundary=np.nan)
      y_overlapped_centers = overlap(y_centers, depth=1, boundary=np.nan)
  
      def chunk(np_zs, np_x_centers, np_y_centers):
  
          # Handle boundaries that have nothing to overlap with
          for centers in [np_x_centers, np_y_centers]:
              if np.isnan(centers[0, :]).all():
                  centers[0, :] = centers[1, :] - (centers[2, :] - centers[1, :])
              if np.isnan(centers[-1, :]).all():
                  centers[-1, :] = centers[-2, :] + (centers[-2, :] - centers[-3, :])
              if np.isnan(centers[:, 0]).all():
                  centers[:, 0] = centers[:, 1] - (centers[:, 2] - centers[:, 1])
              if np.isnan(centers[:, -1]).all():
                  centers[:, -1] = centers[:, -2] + (centers[:, -2] - centers[:, -3])
  
          # compute interval breaks
          x_breaks_chunk = glyph.infer_interval_breaks(np_x_centers)
          y_breaks_chunk = glyph.infer_interval_breaks(np_y_centers)
  
          # trim breaks
          x_breaks_chunk = x_breaks_chunk[1:-1, 1:-1]
          y_breaks_chunk = y_breaks_chunk[1:-1, 1:-1]
  
          # Reconstruct dataset for chunk from numpy array and chunk indices
          chunk_coords = {
              x_coord_name: (x_coord_dim_names, np_x_centers[1:-1, 1:-1]),
              y_coord_name: (y_coord_dim_names, np_y_centers[1:-1, 1:-1]),
          }
          chunk_ds = xr.DataArray(
              np_zs, coords=chunk_coords, dims=data_dim_names, name=var_name
          ).to_dataset()
  
          # Initialize aggregation buffers
          aggs = create(shape)
  
          # Perform aggregation
          extend(aggs, chunk_ds, st, bounds,
                 x_breaks=x_breaks_chunk, y_breaks=y_breaks_chunk)
          return aggs
  
      result_name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
  
      z_keys = [k for row in zs.__dask_keys__() for k in row]
      x_overlap_keys = [k for row in x_overlapped_centers.__dask_keys__() for k in row]
      y_overlap_keys = [k for row in y_overlapped_centers.__dask_keys__() for k in row]
  
      result_keys = [(result_name, i) for i in range(len(z_keys))]
  
      dsk = dict(
          (res_k, (chunk, z_k, x_k, y_k))
          for (res_k, z_k, x_k, y_k) in zip(
              result_keys, z_keys, x_overlap_keys, y_overlap_keys
          )
      )
  
      dsk[result_name] = (
          apply, finalize, [(combine, result_keys)],
          dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
               attrs=dict(x_range=x_range, y_range=y_range))
      )
  
      # Add x/y coord tasks to task graph
      dsk.update(x_overlapped_centers.dask)
      dsk.update(y_overlapped_centers.dask)
  
      return dsk, result_name
  
  
  def dask_xarray_lines(
      glyph: LinesXarrayCommonX, xr_ds: xr.Dataset, schema, canvas, summary,
      *, antialias=False, cuda=False,
  ):
      shape, bounds, st, axis = shape_bounds_st_and_axis(xr_ds, canvas, glyph)
  
      # Compile functions
      create, info, append, combine, finalize, antialias_stage_2, antialias_stage_2_funcs, \
          column_names = compile_components(summary, schema, glyph, antialias=antialias, cuda=cuda,
                                            partitioned=True)
      x_mapper = canvas.x_axis.mapper
      y_mapper = canvas.y_axis.mapper
      extend = glyph._build_extend(
          x_mapper, y_mapper, info, append, antialias_stage_2, antialias_stage_2_funcs)
      x_range = bounds[:2]
      y_range = bounds[2:]
  
      x_name = glyph.x
      x_dim_index = glyph.x_dim_index
      other_dim_index = 1 - x_dim_index
      other_dim_name = xr_ds[glyph.y].coords.dims[other_dim_index]
      xs = xr_ds[x_name]
  
      # Build chunk offsets for coordinates
      chunk_offsets = {}
      for k, chunks in xr_ds.chunks.items():
          chunk_offsets[k] = [0] + list(np.cumsum(chunks))
  
      partitioned = True
      uses_row_index = summary.uses_row_index(cuda, partitioned)
  
      def chunk(np_array, *chunk_indices):
          aggs = create(shape)
  
          start_x_index = chunk_offsets[x_name][chunk_indices[x_dim_index]]
          end_x_index = start_x_index + np_array.shape[x_dim_index]
          x = xs[start_x_index:end_x_index].values
  
          start_other_index = chunk_offsets[other_dim_name][chunk_indices[other_dim_index]]
          end_other_index = start_other_index + np_array.shape[other_dim_index]
  
          data_vars = dict(
              name=(("x", other_dim_name) if x_dim_index == 0 else (other_dim_name, "x"), np_array),
          )
          # Other required columns are chunked in the other_dim
          for column_name in column_names:
              values = xr_ds[column_name][start_other_index:end_other_index].values
              data_vars[column_name] = (other_dim_name, values)
  
          chunk_ds = xr.Dataset(
              data_vars=data_vars,
              coords=dict(
                  x=("x", x),
                  other_dim_name=(other_dim_name, np.arange(start_other_index, end_other_index)),
              ),
          )
  
          if uses_row_index:
              row_offset = start_other_index
              chunk_ds.attrs["_datashader_row_offset"] = row_offset
              chunk_ds.attrs["_datashader_row_length"] = end_other_index - start_other_index
  
          extend(aggs, chunk_ds, st, bounds)
          return aggs
  
      name = tokenize(xr_ds.__dask_tokenize__(), canvas, glyph, summary)
      keys = [k for row in xr_ds.__dask_keys__()[0] for k in row]
      keys2 = [(name, i) for i in range(len(keys))]
      dsk = dict((k2, (chunk, k, k[1], k[2])) for (k2, k) in zip(keys2, keys))
      dsk[name] = (apply, finalize, [(combine, keys2)],
                   dict(cuda=cuda, coords=axis, dims=[glyph.y_label, glyph.x_label],
                        attrs=dict(x_range=x_range, y_range=y_range)))
      return dsk, name
  
  
  dask_glyph_dispatch.register(QuadMeshRectilinear)(dask_rectilinear)
  dask_glyph_dispatch.register(QuadMeshRaster)(dask_raster)
  dask_glyph_dispatch.register(QuadMeshCurvilinear)(dask_curvilinear)
  dask_glyph_dispatch.register(LinesXarrayCommonX)(dask_xarray_lines)
  