<jupyter_start><jupyter_text>Core data model Xarray-Beam tries to make it _straightforward_ to write distributed pipelines with Xarray objects, but unlike libraries like [Xarray with Dask](http://xarray.pydata.org/en/stable/user-guide/dask.html) or Dask/Spark DataFrames, it doesn't hide the distributed magic inside high-level objects.Xarray-Beam is a lower-level tool. You will be manipulating large datasets piece-by-piece yourself, and you as the developer will be responsible for maintaining Xarray-Beam's internal invariants. This means that to successfully use Xarray-Beam, **you will need to understand how how it represents distributed datasets**.This responsibility requires a bit more coding and understanding, but offers benefits in performance and flexibility. This brief tutorial will show you how.We'll start off with some standard imports:<jupyter_code>import apache_beam as beam
import numpy as np
import xarray_beam as xbeam
import xarray<jupyter_output><empty_output><jupyter_text>Keys in Xarray-BeamXarray-Beam is designed around the model that every stage in your Beam pipeline _could_ be stored in a single `xarray.Dataset` object, but is instead represented by a distributed beam `PCollection` of smaller `xarray.Dataset` objects, distributed in two possible ways:- Distinct _variables_ in a Dataset may be separated across multiple records.- Individual arrays can also be split into multiple _chunks_, similar to those used by [dask.array](https://docs.dask.org/en/latest/array.html).To keep track of how individual records could be combined into a larger (virtual) dataset, Xarray-Beam defines a {py:class}`~xarray_beam.Key` object. Key objects consist of:1. `offsets`: integer offests for chunks from the origin in an `immutabledict`2. `vars`: The subset of variables included in each chunk, either as a `frozenset`, or as `None` to indicate "all variables". Making a {py:class}`~xarray_beam.Key` from scratch is simple:<jupyter_code>key = xbeam.Key({'x': 0, 'y': 10}, vars=None)
key<jupyter_output><empty_output><jupyter_text>Or given an existing {py:class}`~xarray_beam.Key`, you can easily modify it with `replace()` or `with_offsets()`:<jupyter_code>key.replace(vars={'foo', 'bar'})
key.with_offsets(x=None, z=1)<jupyter_output><empty_output><jupyter_text>{py:class}`~xarray_beam.Key` objects don't do very much. They are just simple structs with two attributes, along with various special methods required to use them as `dict` keys or as keys in Beam pipelines. You can find a more examples of manipulating keys {py:class}`in its docstring `.  Creating PCollections The standard inputs & outputs for Xarray-Beam are PCollections of `(xbeam.Key, xarray.Dataset)` pairs. Xarray-Beam provides a bunch of PCollections for typical tasks, but many pipelines will still involve some manual manipulation of `Key` and `Dataset` objects, e.g., with builtin Beam transforms like `beam.Map`.To start off, let's write a helper functions for creating our first collection from scratch:<jupyter_code>def create_records():
    for offset in [0, 4]:
        key = xbeam.Key({'x': offset, 'y': 0})
        data = 2 * offset + np.arange(8).reshape(4, 2)
        chunk = xarray.Dataset({
            'foo': (('x', 'y'), data),
            'bar': (('x', 'y'), 100 + data),
        })
        yield key, chunk<jupyter_output><empty_output><jupyter_text>Let's take a look the entries, which are lazily constructed with the generator:<jupyter_code>inputs = list(create_records())
inputs<jupyter_output><empty_output><jupyter_text>```{note}There are multiple valid ways to represent a chunk of a larger dataset with a `Key`.- **Offsets for unchunked dimensions are optional**. Because all chunks have the same offset along the `y` axis, including `y` in `offsets` is not required as long as we don't need to create multiple chunks along that dimension.- **Indicating variables is optional, if all chunks have the same variables**. We could have set `vars={'foo', 'bar'}` on each of these `Key` objects instead of `vars=None`. This would be an equally valid representation of the same records, since all of our datasets have the same variables.``` We now have the inputs we need to use Xarray-Beam's helper functions and PTransforms. For example, we can fully consolidate chunks & variables to see what single `xarray.Dataset` these values would correspond to:<jupyter_code>xbeam.consolidate_fully(inputs)<jupyter_output><empty_output><jupyter_text>To execute with Beam, of course, we need to turn Python lists/generators into Beam PCollections, e.g., with `beam.Create()`:<jupyter_code>with beam.Pipeline() as p:
    p | beam.Create(create_records()) | beam.Map(print)<jupyter_output>WARNING:root:Make sure that locally built Python SDK docker image has Python 3.9 interpreter.<jupyter_text>Writing pipelines Transforms in Xarray-Beam typically act on (key, value) pairs of `(xbeam.Key, xarray.Dataset)`. For example, we can dump our dataset on disk in the scalable [Zarr](https://zarr.readthedocs.io/) format using {py:class}`~xarray_beam.ChunksToZarr`:<jupyter_code>inputs | xbeam.ChunksToZarr('my-data.zarr')<jupyter_output>WARNING:apache_beam.options.pipeline_options:Discarding unparseable args: ['/Users/shoyer/miniconda3/envs/xarray-beam/lib/python3.9/site-packages/ipykernel_launcher.py', '-f', '/Users/shoyer/Library/Jupyter/runtime/kernel-f0f63827-5e6c-42f6-9eda-1f445e1ba489.json']
WARNING:root:Make sure that locally built Python SDK docker image has Python 3.9 interpreter.<jupyter_text>Xarray-Beam doesn't try to provide transformations for everything. In particular, it omits most [embarrassingly parallel](https://en.wikipedia.org/wiki/Embarrassingly_parallel) operations that can be performed independently on each chunk of a larger dataset. You can write these yourself using [`beam.Map`](https://beam.apache.org/documentation/transforms/python/elementwise/map/).For example, consider elementwise arithmetic. We can write a `lambda` function that acts on each key-value pair updating the xarray.Dataset objects appropriately, and put it into an Xarray-Beam pipeline using `beam.MapTuple`:<jupyter_code>inputs | beam.MapTuple(lambda k, v: (k, v + 1))<jupyter_output>WARNING:apache_beam.options.pipeline_options:Discarding unparseable args: ['/Users/shoyer/miniconda3/envs/xarray-beam/lib/python3.9/site-packages/ipykernel_launcher.py', '-f', '/Users/shoyer/Library/Jupyter/runtime/kernel-f0f63827-5e6c-42f6-9eda-1f445e1ba489.json']
WARNING:root:Make sure that locally built Python SDK docker image has Python 3.9 interpreter.<jupyter_text>For operations that add or remove (unchunked) dimensions, you may need to update `Key` objects as well to maintain the Xarray-Beam invariants, e.g., if we want to remove the `y` dimension entirely:<jupyter_code>inputs | beam.MapTuple(lambda k, v: (k.with_offsets(y=None), v.mean('y')))<jupyter_output>WARNING:apache_beam.options.pipeline_options:Discarding unparseable args: ['/Users/shoyer/miniconda3/envs/xarray-beam/lib/python3.9/site-packages/ipykernel_launcher.py', '-f', '/Users/shoyer/Library/Jupyter/runtime/kernel-f0f63827-5e6c-42f6-9eda-1f445e1ba489.json']
WARNING:root:Make sure that locally built Python SDK docker image has Python 3.9 interpreter.
{'fraction_non_alphanumeric': 0.09885089298075592, 'fraction_numerical': 0.013013983109511284, 'mean_word_length': 5.651933701657459, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '565278', 'license': 'Apache-2.0', 'n_tokens_mistral': 2248, 'n_tokens_neox': 2132, 'n_words': 768, 'path': 'docs/data-model.ipynb', 'repo_name': 'alxmrs/xarray-beam', 'size': '7223'}