Skip to content

Commit

Permalink
Merge pull request #154 from dask-contrib/bare_pickle
Browse files Browse the repository at this point in the history
Make unfilled computed hist pickleable
  • Loading branch information
martindurant authored Oct 15, 2024
2 parents 8923899 + ab8e5ac commit f2c91b5
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/dask_histogram/boost.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,13 @@ def __dask_tokenize__(self) -> str:
return self.dask_name

def __dask_postcompute__(self) -> Any:
return lambda x: self._in_memory_type(first(x)), ()
def f(x):
out = self._in_memory_type(first(x))
if hasattr(out, "_dask"):
del out._dask
return out

return f, ()

def __dask_postpersist__(self) -> Any:
return self._rebuild, ()
Expand Down
14 changes: 14 additions & 0 deletions tests/test_boost.py
Original file line number Diff line number Diff line change
Expand Up @@ -553,3 +553,17 @@ def test_histref_pickle():
h1.fill(x) # forces the internal state histref update

pickle.dumps(h1._histref)


def test_boost_output_pickles():
import pickle

import boost_histogram
import dask

import dask_histogram.boost

h = dask_histogram.boost.Histogram(boost_histogram.axis.Regular(10, 0, 1))

o = dask.compute(h)
pickle.dumps(o)

0 comments on commit f2c91b5

Please sign in to comment.