Merge pull request #154 from dask-contrib/bare_pickle

Make unfilled computed hist pickleable

Author: martindurant

src/dask_histogram/boost.py

@@ -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, ()

tests/test_boost.py

@@ -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)