Skip to content

feat: read and write npz methods to replace read_pickle - #1908

Merged
fmaussion merged 79 commits into
OGGM:masterfrom
gampnico:fix-pickle-doomsday
Oct 1, 2026
Merged

fmaussion merged 79 commits into
OGGM:masterfrom
gampnico:fix-pickle-doomsday

Conversation

@gampnico

@gampnico gampnico commented Apr 30, 2026 •

Copy link
Copy Markdown
Member

This PR currently adds:

  • read_store, read_npz and write_npz methods which fall back to pickle if no npz files are available.
  • Warnings that read_pickle methods will be deprecated for future GlacierDirectories.
  • A transcoding module transcoder.py that handles transcoding between existing pickles and the new npz files.
  • Compatibility wrappers in _compat.py, which would be re-used across Feat: incremental glacier directory creation #1951 and feat: stream glacier directories #1966 for reproducible conversion.

NB: This PR's design has changed substantially and has moved to using npz instead of zarr for performance.

Closes: #1903

Points for discussion

All pickles are now stored as npz files in a glacier directory's subdirectory, which I'm currently naming "data_store". Zarr was too inefficient for such small data.

The default replacement for read_pickle is now read_store, which handles both pickle (read_pickle) and numpy (read_npz). read_store falls back to read_pickle if no npz file exists. This preserves backwards compatibility with existing gdirs if a user selects an older URL.

I've designed this as a type registry, rather than mapping to basenames, as it's less brittle, and I don't think it would prevent basename mapping later on. This way we only need to make changes when Python/numpy introduces or deprecates a type (so basically never), rather than when we change or add a new OGGM object.

Zarr does not support all the data structures used by certain pickles (e.g. shapely objects). Previous zarr commits converted these to a compatible type. When using read_store, objects are converted back into the type expected by OGGM. This should minimise rewrites across the rest of the codebase. However, zarr-dependent packages (e.g. DTCG) will require compatibility wrappers.

Closes #1903

  • Tests added/passed
  • Fully documented
  • Entry in whats-new.rst

Keeping this section for posterity, to help explain how the design decisions evolved.

All pickles are now stored as xr.DataTrees in a single zarr file, which I'm currently naming "data_store". The filename in read_zarr(filename) therefore reads the group, rather than entire zarr file. This avoids having multiple small zarr files, and we can maintain cross-compatibility with DTCG.

The default replacement for read_pickle is now read_store, which handles both pickle (read_pickle) and zarr (read_zarr). read_store falls back to read_pickle if no zarr store exists. This preserves backwards compatibility with existing gdirs if a user selects an older URL.

Zarr does not support all the data structures used by certain pickles (e.g. shapely objects). oggm-zarr converts these to a compatible type. When using read_store, these objects are converted back into the type expected by OGGM. This should minimise rewrites across the rest of the codebase.

I'd like to upload a small sample zarr to oggm-sample-data, so this can be included in init_hef, and I can then add tests for reading a zarr from a gdir directly.

I'm currently keeping code for converting from pickle to zarr as a dual-licensed package, as some of this includes code from DTCG which is not compatible with OGGM's license. I'll see if I can rework this, but it may be simpler to handle conversion via dtcg, since this is already set up and we can build directly on the existing GeoZarrHandler.

You can view a conversion workflow here.

Adds:
  - `read_zarr` and `write_zarr` methods which fall back to pickle if no zarr store is available.
  - Warnings that read_pickle methods will be deprecated for future GlacierDirectories.
  - Zarr as a core dependency.

Refs: OGGM#1903
@fmaussion

Copy link
Copy Markdown
Member

Thanks! Looking promising

A few quick thoughts:

  • I don't understand the PR yet - it does not touch the code that needs to be changed, i.e. the places where the pickles are actually written?
  • the read_zarr, read_store and _validate_store are written with the assumption that they are just replacing read_pickle - I think these may become more generic in the future. Maybe we need to rename them according to what they do. Alternatively they could be external functions (not methods of gdir) because they are dealing only with some specific data formats within the oggm workflow - to discuss.
  • I do not want a new oggm-zarr package unless strictly necessary / there are good arguments for it - we can deal with the license like usual, by including the license in OGGM. This could be a module in OGGM if needed
  • It should not be necessary to resort to oggm-sample-data for tests, as OGGM writes and reads pickles (there are no pickles in oggm-sample-data either). In theory, since the change will be mostly internal, the existing tests should "just work" after the conversion to zarr (except those where read_pickle is used), and we need to add a test for backwards compatibility, although I think
    def test_start_from_level_3(self):
    will cover a little bit of this

@gampnico

gampnico commented May 1, 2026

Copy link
Copy Markdown
Member Author

Notes based on internal meeting:

I don't understand the PR yet - it does not touch the code that needs to be changed, i.e. the places where the pickles are actually written?

Next on agenda.

the read_zarr, read_store and _validate_store are written with the assumption that they are just replacing read_pickle - I think these may become more generic in the future. Maybe we need to rename them according to what they do. Alternatively they could be external functions (not methods of gdir) because they are dealing only with some specific data formats within the oggm workflow - to discuss.

Agreed. _validate_store is for converting to and from zarr so that data is structured in a way that OGGM recognises. As discussed we may need separate helper functions for writing specific files (e.g. write_model_flowlines). In pseudo-code:

model_flowlines: Centerline = Centerline(...)
write_store(model_flowlines, "model_flowlines")
model_flowlines: Centerline = read_store("model_flowlines")

where internally read_store and write_store both do something like:

def write_store(self, obj, filename, filesuffix, **kwargs):
    data: xr.DataTree = _validate_store(obj)  # converts shapely objects etc. into zarr-compatible types
    write_zarr(data: xr.DataTree=data, group:str=filename, **kwargs)

I do not want a new oggm-zarr package unless strictly necessary / there are good arguments for it - we can deal with the license like usual, by including the license in OGGM. This could be a module in OGGM if needed

Zarr-related code will be placed under utils as a public module. I'd prefer to keep read_store etc under _workflow as this is where read_pickle is located.

gampnico added 3 commits May 11, 2026 13:03
Adds:
  - Validation for different zarr/pickle stores to maintain backwards
    compatibility.
  - Dedicated module for zarr-related operations.
  - Conversion between existing pickles and zarr.
  - Reconstruction of OGGM objects from minimum possible data.

Refactors:
  - Replaces all instances of `read_pickle` with `read_store`.
    Compatibility is maintained
@gampnico

Copy link
Copy Markdown
Member Author

Failing tests:

@gampnico

Copy link
Copy Markdown
Member Author

@fmaussion this is ready for deployment.

Workaround for datetime64 at ns, which uses an int that can't be decoded
without a unit.
Comment thread pyproject.toml Outdated
"scipy",
"shapely",
"xarray",
"xvec>=0.5.2",

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can be removed

Comment thread docs/whats-new.rst Outdated
``read_store`` and ``write_store``. The latter uses npz files inside glacier
directories instead of pickles. Users should migrate to the new functions, as
compatibility with pickles is not guaranteed in the future. (:pull:`1908`).
By `Nicolas Gampierakis <https://github.com/gampnico>`_.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to be in the 1.X section above now

@fmaussion

fmaussion commented Sep 25, 2026 •

Copy link
Copy Markdown
Member

Excellent work, thanks @gampnico - two minor things in the code, and one last medium thing.

Unless my AI is mistaken, we need two additional small checks:

1. read_pickle fails on directories written with the new store. All library code now uses read_store, but downstream code (tutorials, PyGEM, user scripts) still calls gdir.read_pickle('model_flowlines') etc. On any gdir processed with this branch that raises FileNotFoundError, since only the npz exists.

2. read_pickle can silently return stale data. write_pickle deletes the matching npz, but write_npz doesn't delete the matching .pkl. So if a pickle already exists (e.g. a prepro directory from the server) and a task rewrites the data via write_store, read_store gets the new data but read_pickle quietly gets the old one:

fls = gdir.read_store('model_flowlines')
gdir.write_pickle(fls[:1], 'model_flowlines', filesuffix='_test')  # pre-existing pickle
gdir.write_store(fls, 'model_flowlines', filesuffix='_test')       # npz written, .pkl kept
len(gdir.read_store('model_flowlines', '_test'))                   # 3
len(gdir.read_pickle('model_flowlines', filesuffix='_test'))       # 1  <- stale

The missing file at least raises an error; the stale pickle is the issue

Suggested fix, mirroring what write_pickle already does:

  • In write_npz, after the os.replace, remove self.get_filepath(filename, filesuffix=filesuffix) if it exists, so there is only ever one copy on disk.
  • In read_pickle, if the .pkl doesn't exist but the npz does, return self.read_npz(...) (maybe with a DeprecationWarning pointing to read_store). There's no recursion risk: read_store only falls back to read_pickle when the .pkl exists.

convert_pickles_to_npz(delete=False) would need a tweak, since write_store would now remove the pickle itself. It's probably simplest to drop the delete option, since keeping both copies is exactly what we want to avoid.

A regression test for the snippet above would be good too.

Almost ready!

@gampnico

gampnico commented Sep 28, 2026 •

Copy link
Copy Markdown
Member Author

For point 1:

read_pickle fails on directories written with the new store

This is intentional behaviour, and I've already opened a PR for the tutorials: OGGM/tutorials#208 . read_pickle cannot read a pickle if no such pickle exists. I removed my previous deprecation warnings because it absolutely tanks test performance on multiple read/writes, but I can add this back, as the impact was only on tests.

If you like I could add a message with FileNotFoundError telling the user to update from read_pickle to read_store?

We can maybe chat about point 2 later today. If the default is _store, then the user won't read the stale pickle unless they explicitly call read_pickle. I delete the pickles by default when using convert_pickles_to_npz in _compat.py:

if os.path.isfile(npz_fp) and delete:
    # npz write succeeded, drop the now-redundant pickle
    os.remove(fp)

But this is done with an explicit argument, and I'm not sure if it's safe to do so silently. What do you think?

In write_npz, after the os.replace, remove self.get_filepath(filename, filesuffix=filesuffix) if it exists, so there is only ever one copy on disk.

I'm not sure what this is referring to, write_npz doesn't call self.get_filepath. Are you referring to the call in read_store? If so, this is the fallback mechanism that looks for a pickle if no npz is found.

My interpretation with read_pickle is this is part of a deprecation cycle, so we wouldn't be supporting read_pickle, as the concept of "pickles" in OGGM and GlacierDirectory would no longer exist. I would argue that keeping _pickle makes it more likely for users to read a stale pickle, because they're not using _store to manage this automatically.

This depends on how you want this phased out, and where you want to introduce a breaking change in OGGM's release cycle. If you want something like:

  1. fix-pickle-doomsday: _pickle falls back to _npz and issues a DeprecationWarning. Users can use _pickle, but potentially don't understand why they need to switch to _store.
  2. feat-ON-94-incremental-gdirs: _pickle become private methods for fallback mechanisms or no longer work. A user can still explicitly call _read_pickle but it's at their own risk and we no longer support it. Manifest writes etc. won't work properly.
  3. feat-ON-95-streaming-gdirs: we have to remove pickles because they're incompatible.

@fmaussion

Copy link
Copy Markdown
Member

Good points! Thanks Nic. I guess my first reaction when I thought about it was: "Since the objects in the background didn't change, why not keep read_pickle working". Thinking about PyGEM and other "wild" OGGM code in particular, the change to npz is just a small annoyance for them. But PyGEM uses read_pickle only 9 times: https://github.com/search?q=repo%3APyGEM-Community%2FPyGEM+read_pickle&type=code

So (1) and (2) could be solved in one go if we say read_pickle simply is meant to disappear: it errors if any npz file is available in the dir, and write_pickle can simply disappear entirely? We can discuss later today.

@gampnico

Copy link
Copy Markdown
Member Author

Based on our discussion today, read_pickle and write_pickle are now private. This is a breaking change that should be clearly indicated to users, who will have to update their own code accordingly. I've kept them in workflow.py as they are still part of read/write_store fallback mechanisms.

@fmaussion

Copy link
Copy Markdown
Member

We are good to go here, right?

@gampnico

Copy link
Copy Markdown
Member Author

Yes, all good to go!

@gampnico

Copy link
Copy Markdown
Member Author

Apologies for the accidental push, this is now ready to merge.

@fmaussion
fmaussion merged commit e62389d into OGGM:master Oct 1, 2026
30 checks passed
@fmaussion

Copy link
Copy Markdown
Member

Thanks @gampnico ! Great to see this in

fmaussion pushed a commit to fmaussion/oggm that referenced this pull request Oct 3, 2026
fmaussion added a commit that referenced this pull request Oct 4, 2026
* allow prcp fac fixed with new temp_bias system

* revert unnecessary changes in PR #1908

---------

Co-authored-by: Fabien Maussion <Fabien Maussion>
@gampnico
gampnico deleted the fix-pickle-doomsday branch November 3, 2026 11:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Get rid of pickles, one pickle at a time

2 participants