Adding chunk & type information to dask high level graphs

dask/array/core.py

@@ -1156,6 +1156,24 @@ def __new__(cls, dask, name, chunks, dtype=None, meta=None, shape=None):
             if result is not None:
                 self = result
 
+        try:
+            layer = self.dask.layers[name]
+        except (AttributeError, KeyError):
+            # self is no longer an Array after applying the plugins, OR
+            # a plugin replaced the HighLevelGraph with a plain dict, OR
+            # name is not the top layer's name (this can happen after the layer is
+            # manipulated, to avoid a collision)
+            pass
+        else:
+            layer.collection_annotations.update(
+                {
+                    "type": type(self),
+                    "chunk_type": type(meta),
+                    "chunks": chunks,
+                    "dtype": dtype,
+                }
+            )
+
         return self
 
     def __reduce__(self):

dask/dataframe/core.py

@@ -3609,6 +3609,20 @@ class DataFrame(_Frame):
     _token_prefix = "dataframe-"
     _accessors = set()
 
+    def __init__(self, dsk, name, meta, divisions):
+        super().__init__(dsk, name, meta, divisions)
+        self.dask.layers[name].collection_annotations.update(
+            {
+                "type": type(self),
+                "divisions": divisions,
+                "dataframe_type": type(meta),
+                "series_dtypes": {
+                    col: meta[col].dtype if hasattr(meta[col], "dtype") else None
+                    for col in meta.columns
+                },
+            }
+        )
+
     def __array_wrap__(self, array, context=None):
         if isinstance(context, tuple) and len(context) > 0:
             if isinstance(context[1][0], np.ndarray) and context[1][0].shape == ():

dask/highlevelgraph.py

@@ -57,12 +57,40 @@ class Layer(collections.abc.Mapping):
     """
 
     annotations: Optional[Mapping[str, Any]]
+    collection_annotations: Optional[Mapping[str, Any]]
 
-    def __init__(self, annotations: Mapping[str, Any] = None):
+    def __init__(
+        self,
+        annotations: Mapping[str, Any] = None,
+        collection_annotations: Mapping[str, Any] = {},
+    ):
+        """Initialize Layer object.
+
+        Parameters
+        ----------
+        annotations : Mapping[str, Any], optional
+            By default, None.
+            Annotations are metadata or soft constraints associated with tasks
+            that dask schedulers may choose to respect:
+            They signal intent without enforcing hard constraints.
+            As such, they are primarily designed for use with the distributed
+            scheduler. See the dask.annotate function for more information.
+        collection_annotations : Mapping[str, Any], optional
+            By default an empty dict.
+            Experimental, intended to assist with visualizing the performance
+            characteristics of Dask computations.
+            These annotations are *not* passed to the distributed scheduler.
+        """
         if annotations:
             self.annotations = annotations
         else:
             self.annotations = copy.copy(config.get("annotations", None))
+        if collection_annotations:
+            self.collection_annotations = collection_annotations
+        else:
+            self.collection_annotations = copy.copy(
+                config.get("collection_annotations", {})
+            )
 
     @abc.abstractmethod
     def is_materialized(self) -> bool: