allow renaming when rebuilding

pfackeldey · Dec 3, 2024 · b375778 · b375778
1 parent 1d4d4e9
commit b375778
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py
@@ -381,7 +381,7 @@ def __dask_postpersist__(self):
     def _rebuild(self, dsk, *, rename=None):
         name = self._name
         if rename:
-            raise ValueError("rename= unsupported in dask-awkward")
+            name = rename.get(name, name)
         return type(self)(dsk, name, self._meta, self.known_value)
 
     def __reduce__(self):
@@ -962,7 +962,7 @@ def __setitem__(self, where: Any, what: Any) -> None:
     def _rebuild(self, dsk, *, rename=None):
         name = self.name
         if rename:
-            raise ValueError("rename= unsupported in dask-awkward")
+            name = rename.get(name, name)
         return Array(dsk, name, self._meta, divisions=self.divisions)
 
     def reset_meta(self) -> None: