Add rechunking to save_dataset

[benjeffery]

I'm creating a workflow where I need to subset the variants in a dataset and write it back to zarr. However:

ds = sg.simulate_genotype_call_dataset(50000,10)
sg.save_dataset('test.zarr')
ds = sg.load_dataset('test.zarr')
selection = xr.ones_like(ds['variant_position'], dtype=bool)
selection[1000:1100] = False
ds = ds.sel(variants=selection)
sg.save_dataset('test.zarr')

Results in: ValueError: Final chunk of Zarr array must be the same size or smaller than the first.
Or with more chunks: ValueError: Zarr requires uniform chunk sizes except for final chunk.

I could rechunk all the arrays with a variants dimension - this doesn't make for an easy time for the average user though.
Should save_dataset auto rechunk in these situations?

[tomwhite]

Good question. When we added save_dataset we didn't include auto-rechunking, since it can perform poorly depending on chunk size, and it's good for the user to be aware of what's happening.

At a minimum we could improve the documentation, and give some examples (including at https://pystatgen.github.io/sgkit/latest/how_do_i.html#save-results-to-a-zarr-file).

We could also add an optional parameter to save_dataset to control rechunking; perhaps opt-in?

[benjeffery]

Thanks @tomwhite. I agree with your point about auto-rechunking being a heavy operation and it should be opt-in. I think adding an opt-in parameter to save_dataset would be nice. save_dataset could then detect that re-chunking was needed and throw a friendly error message pointing out the flag?

[tomwhite]

I agree - this sounds like a good addition.

[tomwhite]

This was mentioned in #296 too.