Switch Index to be a frozenset

[rabernat]

Simplification suggested by Stephan in google/xarray-beam#32 (comment).

Update:

I was amazed how simple it was to make this change. I just swapped our our Index class with

Index = frozenset[DimIndex]

and mostly everything just worked, requiring no change to how indexes are created (e.g. idx = Index((some, stuff))). That was a nice surprise.

However, this PR has turned up some nasty mypy-related stuff, and also showed that that simple approach doesn't work in 3.8. In 3.8, we would need to do something like

Index = FrozenSet[DimIndex]

However, if we do that, we get the error TypeError: Type FrozenSet cannot be instantiated; use frozenset() instead when we try to instantiate the Index.

I would like to say we will just drop python 3.8. However, afaik, beam does not yet work with 3.9. So it looks like we will need some workarounds.

Does anyone know the best way to solve this?

[rabernat]

Tests have shown that this does not work on python 3.8, even with from __future__ import annotations

[rabernat]

Ok I think I found the solution:

class Index(FrozenSet[DimIndex]):
    pass

[rabernat]

The upstream-dev tests have surfaced a new issue

        # https://github.com/pangeo-forge/pangeo-forge-recipes/issues/151
        # Requires >4 GiB of memory to run.
        pd = pytest.importorskip("pandas")
        distributed = pytest.importorskip("distributed")
    
        dates = pd.date_range("2020-05-31T00:00:00", "2021-05-31T23:59:59", freq="30min")
        time_concat_dim = ConcatDim("time", dates, nitems_per_file=1)
        pattern = FilePattern(_make_filename_for_memory_usage_test, time_concat_dim,)
    
        recipe = XarrayZarrRecipe(
            pattern, xarray_open_kwargs={"group": "Grid", "decode_coords": "all"}, inputs_per_chunk=1,
        )
    
>       delayed = recipe.to_dask()

tests/recipe_tests/test_graph_memory.py:49: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
pangeo_forge_recipes/recipes/base.py:28: in to_dask
    return DaskPipelineExecutor.compile(self._compiler())
pangeo_forge_recipes/executors/dask.py:75: in compile
    delayed = Delayed(prev_token[0], hlg)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = Delayed('finalize_target-2f31500925ec7e187d227c54ab925c75')
key = 'finalize_target-2f31500925ec7e187d227c54ab925c75'
dsk = HighLevelGraph with 5 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x7f63f0f6e2b0>
 0. config
 1. cache_input
 2. prepare_target
 3. store_chunk
 4. finalize_target

length = None, layer = None

    def __init__(self, key, dsk, length=None, layer=None):
        self._key = key
        self._dask = dsk
        self._length = length
    
        # NOTE: Layer is used by `to_delayed` in other collections, but not in normal Delayed use
        self._layer = layer or key
        if isinstance(dsk, HighLevelGraph) and self._layer not in dsk.layers:
>           raise ValueError(
                f"Layer {self._layer} not in the HighLevelGraph's layers: {list(dsk.layers)}"
            )
E           ValueError: Layer finalize_target-2f31500925ec7e187d227c54ab925c75 not in the HighLevelGraph's layers: ['config', 'cache_input', 'prepare_target', 'store_chunk', 'finalize_target']

Apparently the way we are constructing high-level graphs is not compatible with the latest dask release.

I'm inclined to merge this, let the master tests fail, and open a new issue for it.

[rabernat]

/run-test-tutorials

[cisaacstern]

I'm inclined to merge this, let the master tests fail, and open a new issue for it.

👍 upstream-dev tests ftw

[cisaacstern]

Out of curiosity, is there a one-ish sentence way to encapsulate the advantage of a frozenset vs. a tuple here? I wasn't able to immediately grasp the motivation from the linked xarray-beam thread.

[rabernat]

This class was already basically doing the same thing as a frozenset! So this change simplifies our Index type--a key data structure--to use a built-in type. Simpler and less code to maintain.

For context, this is a tiny part of the changes that will move us in the direction of #256.