get rid of hard coded keys

Author: fjetter

dask_expr/_concat.py

@@ -271,7 +271,7 @@ class StackPartition(Concat):
     _defaults = {"join": "outer", "ignore_order": False, "_kwargs": {}, "axis": 0}
 
     def _layer(self):
-        dsk, i = {}, 0
+        dsk = {}
         kwargs = self._kwargs.copy()
         kwargs["ignore_order"] = self.ignore_order
         ctr = 0
@@ -282,16 +282,15 @@ def _layer(self):
                 match = True
             except (ValueError, TypeError):
                 match = False
-
-            for i in range(df.npartitions):
+            for key in df.__dask_keys__():
                 if match:
-                    dsk[(self._name, ctr)] = df._name, i
+                    dsk[(self._name, ctr)] = key
                 else:
                     dsk[(self._name, ctr)] = (
                         apply,
                         methods.concat,
                         [
-                            [meta, (df._name, i)],
+                            [meta, key],
                             self.axis,
                             self.join,
                             False,
@@ -321,7 +320,7 @@ def _layer(self):
                 apply,
                 methods.concat,
                 [
-                    [(df._name, i) for df in dfs],
+                    [df.__dask_keys__()[i] for df in dfs],
                     self.axis,
                     self.join,
                     False,

dask_expr/_cumulative.py

@@ -76,23 +76,25 @@ def _meta(self):
     def _layer(self) -> dict:
         dsk = {}
         frame, previous_partitions = self.frame, self.previous_partitions
-        dsk[(self._name, 0)] = (frame._name, 0)
+        frame_keys = frame.__dask_keys__()
+        previous_partitions_keys = previous_partitions.__dask_keys__()
+        dsk[(self._name, 0)] = frame_keys[0]
 
         intermediate_name = self._name + "-intermediate"
         for i in range(1, self.frame.npartitions):
             if i == 1:
-                dsk[(intermediate_name, i)] = (previous_partitions._name, i - 1)
+                dsk[(intermediate_name, i)] = previous_partitions_keys[i - 1]
             else:
                 # aggregate with previous cumulation results
                 dsk[(intermediate_name, i)] = (
                     methods._cum_aggregate_apply,
                     self.aggregator,
                     (intermediate_name, i - 1),
-                    (previous_partitions._name, i - 1),
+                    previous_partitions_keys[i - 1],
                 )
             dsk[(self._name, i)] = (
                 self.aggregator,
-                (self.frame._name, i),
+                frame_keys[i],
                 (intermediate_name, i),
             )
         return dsk

dask_expr/_expr.py

@@ -517,9 +517,9 @@ def _blockwise_arg(self, arg, i):
         if isinstance(arg, Expr):
             # Make key for Expr-based argument
             if self._broadcast_dep(arg):
-                return (arg._name, 0)
+                return arg.__dask_keys__()[0]
             else:
-                return (arg._name, i)
+                return arg.__dask_keys__()[i]
 
         else:
             return arg
@@ -636,11 +636,15 @@ def _get_meta_ufunc(dfs, args, func):
         raise NotImplementedError(msg)
     # For broadcastable series, use no rows.
     parts = [
-        d._meta
-        if d.ndim == 0
-        else np.empty((), dtype=d.dtype)
-        if isinstance(d, Array)
-        else meta_nonempty(d._meta)
+        (
+            d._meta
+            if d.ndim == 0
+            else (
+                np.empty((), dtype=d.dtype)
+                if isinstance(d, Array)
+                else meta_nonempty(d._meta)
+            )
+        )
         for d in dasks
     ]
 
@@ -895,12 +899,13 @@ def _layer(self) -> dict:
         dsk, prevs, nexts = {}, [], []
 
         name_prepend = "overlap-prepend" + self.frame._name
+        frame_keys = self.frame.__dask_keys__()
         if self.before:
             prevs.append(None)
             if isinstance(self.before, numbers.Integral):
                 before = self.before
                 for i in range(self.frame.npartitions - 1):
-                    dsk[(name_prepend, i)] = (M.tail, (self.frame._name, i), before)
+                    dsk[(name_prepend, i)] = (M.tail, frame_keys[i], before)
                     prevs.append((name_prepend, i))
             elif isinstance(self.before, datetime.timedelta):
                 # Assumes monotonic (increasing?) index
@@ -928,17 +933,17 @@ def _layer(self) -> dict:
 
                         dsk[(name_prepend, i)] = (
                             _tail_timedelta,
-                            (self.frame._name, i + 1),
-                            [(self.frame._name, k) for k in range(j, i + 1)],
+                            frame_keys[i + 1],
+                            [frame_keys[k] for k in range(j, i + 1)],
                             self.before,
                         )
                         prevs.append((name_prepend, i))
                 else:
                     for i in range(self.frame.npartitions - 1):
                         dsk[(name_prepend, i)] = (
                             _tail_timedelta,
-                            (self.frame._name, i + 1),
-                            [(self.frame._name, i)],
+                            frame_keys[i + 1],
+                            [frame_keys[i]],
                             self.before,
                         )
                         prevs.append((name_prepend, i))
@@ -950,7 +955,7 @@ def _layer(self) -> dict:
             if isinstance(self.after, numbers.Integral):
                 after = self.after
                 for i in range(1, self.frame.npartitions):
-                    dsk[(name_append, i)] = (M.head, (self.frame._name, i), after)
+                    dsk[(name_append, i)] = (M.head, frame_keys[i], after)
                     nexts.append((name_append, i))
             else:
                 # We don't want to look at the divisions, so take twice the step and
@@ -959,8 +964,8 @@ def _layer(self) -> dict:
                 for i in range(1, self.frame.npartitions):
                     dsk[(name_append, i)] = (
                         _head_timedelta,
-                        (self.frame._name, i - 1),
-                        (self.frame._name, i),
+                        frame_keys[i - 1],
+                        frame_keys[i],
                         after,
                     )
                     nexts.append((name_append, i))
@@ -974,7 +979,7 @@ def _layer(self) -> dict:
             dsk[(self._name, i)] = (
                 _combined_parts,
                 prev,
-                (self.frame._name, i),
+                frame_keys[i],
                 next,
                 self.before,
                 self.after,
@@ -1605,7 +1610,7 @@ def _task(self, index: int):
             apply,
             M.apply,
             [
-                (self.frame._name, index),
+                self.frame.__dask_keys__()[index],
                 self.function,
             ]
             + list(self.args),
@@ -1825,9 +1830,11 @@ def _meta(self):
         caselist = [
             (
                 meta_nonempty(c[i]._meta) if isinstance(c[i], Expr) else c[i],
-                meta_nonempty(c[i + 1]._meta)
-                if isinstance(c[i + 1], Expr)
-                else c[i + 1],
+                (
+                    meta_nonempty(c[i + 1]._meta)
+                    if isinstance(c[i + 1], Expr)
+                    else c[i + 1]
+                ),
             )
             for i in range(0, len(c), 2)
         ]
@@ -1939,7 +1946,7 @@ def _projection_columns(self):
     def _task(self, index: int):
         return (
             getattr,
-            (self.frame._name, index),
+            self.frame.__dask_keys__()[index],
             "index",
         )
 
@@ -1981,15 +1988,13 @@ def _divisions(self):
     def _layer(self):
         non_empties = [i for i, length in enumerate(self.lens) if length != 0]
         # If all empty, collapse into one partition
+        frame_keys = self.frame.__dask_keys__()
         if len(non_empties) == 0:
-            return {(self._name, 0): (self.frame._name, 0)}
+            return {(self._name, 0): frame_keys[0]}
 
         # drop empty partitions by mapping each partition in a new graph to a particular
         # partition on the old graph.
-        dsk = {
-            (self._name, i): (self.frame._name, div)
-            for i, div in enumerate(non_empties)
-        }
+        dsk = {(self._name, i): frame_keys[div] for i, div in enumerate(non_empties)}
         ddf_keys = list(dsk.values())
 
         overlap = [
@@ -2044,9 +2049,9 @@ def _simplify_down(self):
 
     def _layer(self):
         name = "part-" + self._name
+
         dsk = {
-            (name, i): (len, (self.frame._name, i))
-            for i in range(self.frame.npartitions)
+            (name, i): (len, key) for i, key in enumerate(self.frame.__dask_keys__())
         }
         dsk[(self._name, 0)] = (tuple, list(dsk.keys()))
         return dsk
@@ -2226,12 +2231,12 @@ def _task(self, index: int):
             op = safe_head
         else:
             op = M.head
-        return (op, (self.frame._name, index), self.n)
+        return (op, self.frame.__dask_keys__()[index], self.n)
 
 
 class BlockwiseHeadIndex(BlockwiseHead):
     def _task(self, index: int):
-        return (operator.getitem, (self.frame._name, index), slice(0, self.n))
+        return (operator.getitem, self.frame.__dask_keys__()[index], slice(0, self.n))
 
 
 class Tail(Expr):
@@ -2289,12 +2294,16 @@ def _divisions(self):
         return self.frame.divisions
 
     def _task(self, index: int):
-        return (M.tail, (self.frame._name, index), self.n)
+        return (M.tail, self.frame.__dask_keys__()[index], self.n)
 
 
 class BlockwiseTailIndex(BlockwiseTail):
     def _task(self, index: int):
-        return (operator.getitem, (self.frame._name, index), slice(-self.n, None))
+        return (
+            operator.getitem,
+            self.frame.__dask_keys__()[index],
+            slice(-self.n, None),
+        )
 
 
 class Binop(Elemwise):
@@ -2582,16 +2591,18 @@ def _divisions(self):
         return tuple(divisions)
 
     def _task(self, index: int):
-        return (self.frame._name, self.partitions[index])
+        return self.frame.__dask_keys__()[self.partitions[index]]
 
     def _simplify_down(self):
         if isinstance(self.frame, Blockwise) and not isinstance(
             self.frame, (BlockwiseIO, Fused)
         ):
             operands = [
-                Partitions(op, self.partitions)
-                if (isinstance(op, Expr) and not self.frame._broadcast_dep(op))
-                else op
+                (
+                    Partitions(op, self.partitions)
+                    if (isinstance(op, Expr) and not self.frame._broadcast_dep(op))
+                    else op
+                )
                 for op in self.frame.operands
             ]
             return type(self.frame)(*operands)
@@ -3130,18 +3141,19 @@ def _broadcast_dep(self, dep: Expr):
         return dep.npartitions == 1
 
     def _task(self, index):
-        graph = {self._name: (self.exprs[0]._name, index)}
+        graph = {self._name: self.exprs[0].__dask_keys__()[index]}
         for _expr in self.exprs:
+            frame_keys = _expr.__dask_keys__()
             if isinstance(_expr, Fused):
                 subgraph, name = _expr._task(index)[1:3]
                 graph.update(subgraph)
                 graph[(name, index)] = name
             elif self._broadcast_dep(_expr):
                 # When _expr is being broadcasted, we only
                 # want to define a fused task for index 0
-                graph[(_expr._name, 0)] = _expr._task(0)
+                graph[frame_keys[0]] = _expr._task(0)
             else:
-                graph[(_expr._name, index)] = _expr._task(index)
+                graph[frame_keys[index]] = _expr._task(index)
 
         for i, dep in enumerate(self.dependencies()):
             graph[self._blockwise_arg(dep, index)] = "_" + str(i)
@@ -3247,9 +3259,11 @@ def maybe_align_partitions(*exprs, divisions):
     from dask_expr._repartition import Repartition
 
     return [
-        Repartition(df, new_divisions=divisions, force=True)
-        if isinstance(df, Expr) and df.ndim > 0
-        else df
+        (
+            Repartition(df, new_divisions=divisions, force=True)
+            if isinstance(df, Expr) and df.ndim > 0
+            else df
+        )
         for df in exprs
     ]
 

dask_expr/_groupby.py

@@ -1316,25 +1316,27 @@ def _divisions(self):
         return self.frame.divisions
 
     def _layer(self) -> dict:
-        dsk = {(self._name, 0): (self.cum_raw._name, 0)}
+        dsk = {(self._name, 0): self.cum_raw.__dask_keys__()[0]}
         name_cum = "cum-last" + self._name
-
-        for i in range(1, self.frame.npartitions):
+        cum_last_keys = self.cum_last.__dask_keys__()
+        for i, frame_key in enumerate(self.frame.__dask_keys__()):
+            if i == 0:
+                continue
             # store each cumulative step to graph to reduce computation
             if i == 1:
-                dsk[(name_cum, i)] = (self.cum_last._name, i - 1)
+                dsk[(name_cum, i)] = cum_last_keys[i - 1]
             else:
                 # aggregate with previous cumulation results
                 dsk[(name_cum, i)] = (
                     _cum_agg_filled,
                     (name_cum, i - 1),
-                    (self.cum_last._name, i - 1),
+                    cum_last_keys[i - 1],
                     self.aggregate,
                     self.initial,
                 )
             dsk[(self._name, i)] = (
                 _cum_agg_aligned,
-                (self.frame._name, i),
+                frame_key,
                 (name_cum, i),
                 self.by,
                 self.operand("columns"),