O(1) access to keys in HighLevelGraph

crusaderky · crusaderky · commit 0ad5fe3e029d · 2021-01-21T14:04:59.000Z
diff --git a/Untitled.ipynb b/Untitled.ipynb
diff --git a/dask/blockwise.py b/dask/blockwise.py
@@ -59,8 +59,8 @@ def blockwise(
 ):
     """Create a Blockwise symbolic mutable mapping
 
-    This is like the ``make_blockwise_graph`` function, but rather than construct a dict, it
-    returns a symbolic Blockwise object.
+    This is like the ``make_blockwise_graph`` function, but rather than construct a
+    dict, it returns a symbolic Blockwise object.
 
     See Also
     --------
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
@@ -476,9 +476,27 @@ def from_collections(cls, name, layer, dependencies=()):
         return cls(layers, deps)
 
     def __getitem__(self, key):
+        # Attempt O(1) direct access first, under the assumption that layer names match
+        # either the keys (Delayed, in most cases) or the first element of the key
+        # tuples (Array, Bag, DataFrame, Series). This assumption could be wrong for
+        # third-party collections and/or transforms, and is sometimes wrong for Delayed
+        # objects too.
+        try:
+            return self.layers[key][key]
+        except KeyError:
+            pass
+        try:
+            return self.layers[key[0]][key]
+        except (KeyError, IndexError, TypeError):
+            pass
+
+        # Fall back to O(n) access
         for d in self.layers.values():
-            if key in d:
+            try:
                 return d[key]
+            except KeyError:
+                pass
+
         raise KeyError(key)
 
     def __len__(self):
