Repository: dask/dask

Issue:
Displaying Array/DataFrame in notebook materializes all graph layers
Displaying the `_repr_html_` of an Array or DataFrame includes the number of tasks. Figuring out the number of tasks causes all layers in the HighLevelGraph to be materialized, which makes displaying the repr very slow for large graphs, plus somewhat undoes all the good work of getting more things into Blockwise.

`_repr_html_` gets the task count from `len(self.dask)`. `HighLevelGraph.__len__`, in turn, just returns `len(self.to_dict())`

**Minimal Complete Verifiable Example**:

```python
import dask.array as da
def layers_are_materialized(obj):
    return {k: lyr.is_materialized() for k, lyr in obj.__dask_graph__().layers.items()}

x = da.ones((10000, 10000), chunks=10)

# add some blockwise layers
for i in range(4):
    x = x + 1

layers_are_materialized(x)
# {'ones-e532400d218fbe77c0befc8551fa738a': True,
#  'add-90a3aa343fa08babeff712d337a8f287': False,
#  'add-7c52aeffa2899a67494f584b1b6ce24c': False,
#  'add-ae7ff0a503e9c7464de58300616c3649': False,
#  'add-96feedc695ac0a342339674fb0b1067f': False}

%time _ = x._repr_html_()
CPU times: user 18.8 s, sys: 469 ms, total: 19.3 s
Wall time: 19.3 s

layers_are_materialized(x)
# {'ones-e532400d218fbe77c0befc8551fa738a': True,
#  'add-90a3aa343fa08babeff712d337a8f287': True,
#  'add-7c52aeffa2899a67494f584b1b6ce24c': True,
#  'add-ae7ff0a503e9c7464de58300616c3649': True,
#  'add-96feedc695ac0a342339674fb0b1067f': True}
```

<details><summary>(Just to confirm I profiled, and it is indeed <code>HighLevelGraph.__len__</code> that takes all the time)</summary>

![Screen Shot 2021-02-24 at 6 09 49 PM](https://user-images.githubusercontent.com/3309802/109087870-84008f80-76cb-11eb-94dd-a73869158e76.png)

</details>

**Anything else we need to know?**:

I wasn't sure where to address this—change `__len__` on HighLevelGraph to not materialize the full dict but give an approximation, or have `_repr_html_` use different logic (since it would be okay for display purposes to have the task count be less accurate)?

These are the alternatives I'm imagining:

```python
# sum the lengths of the layers
%%timeit
sum(len(layer) for layer in x.dask.layers.values())
# 11.5 µs ± 244 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

# use get_all_external_keys
%%timeit
len(x.dask.get_all_external_keys())
delattr(x.dask, "_all_external_keys")
# 2.81 s ± 42.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

The first would over-count duplicate keys between layers; the second would under-count intermediate keys. So probably neither is appropriate for `__len__`, but which is more likely to be accurate for `_repr_html_`?

cc @crusaderky since you worked on this most recently.

**Environment**:

- Dask version: master (895c9541044e16809151f0e3db95c48ea338c6d3)
- Python version: 3.8.7


Discussion:
> ``__len__`` on HighLevelGraph to not materialize the full dict

this is feasible, except that it would cause len(hlg) > len(hlg.to_dict()) in the broken case where there are key collisions. Not sure if we care about it; probably not. I'm personally in favor.

> The first would over-count duplicate keys between layers; the second would under-count intermediate keys

duplicate keys between layers is a broken use case; having intermediate keys hidden inside a layer is a legit one, so I'm definitely against under-counting.
Thanks for the info. Duplicate keys between layers didn't make sense to me either, but I didn't know it was actually a broken state. Do you know what can cause this to happen?

Anyway, all this points to `sum(len(layer) for layer in self.layers.values())` being an appropriate implementation for `__len__` on HighLevelGraph. Happy to make that change if you'd like.
> Do you know what can cause this to happen?

1. handcrafted keys and layer names introduce an accidental collision (it's a bug)
2. tokenize() is accidentally called without passing all independent variables into it (it's a bug), AND for some reason you get different results for the layer names but a collision for a key inside the layer; in other words key[0] != layer name (meaning most of the times an internal key)
3. ?

> Happy to make that change if you'd like.

+1 from me

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -270,20 +270,13 @@ def __getitem__(self, key):
     def __iter__(self):
         return iter(self._dict)
 
-    def __len__(self):
-        return int(prod(self._out_numblocks().values()))
-
-    def _out_numblocks(self):
-        d = {}
-        out_d = {}
-        indices = {k: v for k, v in self.indices if v is not None}
-        for k, v in self.numblocks.items():
-            for a, b in zip(indices[k], v):
-                d[a] = max(d.get(a, 0), b)
-                if a in self.output_indices:
-                    out_d[a] = d[a]
-
-        return out_d
+    def __len__(self) -> int:
+        # same method as `get_output_keys`, without manifesting the keys themselves
+        return (
+            len(self.output_blocks)
+            if self.output_blocks
+            else prod(self.dims[i] for i in self.output_indices)
+        )
 
     def is_materialized(self):
         return hasattr(self, "_cached_dict")
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -659,8 +659,13 @@ def __getitem__(self, key):
 
         raise KeyError(key)
 
-    def __len__(self):
-        return len(self.to_dict())
+    def __len__(self) -> int:
+        # NOTE: this will double-count keys that are duplicated between layers, so it's
+        # possible that `len(hlg) > len(hlg.to_dict())`. However, duplicate keys should
+        # not occur through normal use, and their existence would usually be a bug.
+        # So we ignore this case in favor of better performance.
+        # https://github.com/dask/dask/issues/7271
+        return sum(len(layer) for layer in self.layers.values())
 
     def __iter__(self):
         return iter(self.to_dict())


Test patch:
diff --git a/dask/tests/test_highgraph.py b/dask/tests/test_highgraph.py
--- a/dask/tests/test_highgraph.py
+++ b/dask/tests/test_highgraph.py
@@ -1,3 +1,4 @@
+from dask.blockwise import Blockwise, blockwise_token
 import os
 from collections.abc import Set
 
@@ -187,3 +188,27 @@ def test_highlevelgraph_dicts_deprecation():
         layers = {"a": MaterializedLayer({"x": 1, "y": (inc, "x")})}
         hg = HighLevelGraph(layers, {"a": set()})
         assert hg.dicts == layers
+
+
+def test_len_does_not_materialize():
+    a = {"x": 1}
+    b = Blockwise(
+        output="b",
+        output_indices=tuple("ij"),
+        dsk={"b": [[blockwise_token(0)]]},
+        indices=(),
+        numblocks={},
+        new_axes={"i": (1, 1, 1), "j": (1, 1)},
+    )
+    assert len(b) == len(b.get_output_keys())
+
+    layers = {"a": a, "b": b}
+    dependencies = {"a": set(), "b": {"a"}}
+    hg = HighLevelGraph(layers, dependencies)
+
+    assert hg.layers["a"].is_materialized()
+    assert not hg.layers["b"].is_materialized()
+
+    assert len(hg) == len(a) + len(b) == 7
+
+    assert not hg.layers["b"].is_materialized()