Repartition source frame with update_ensemble

src/tape/ensemble.py

@@ -336,7 +336,7 @@ def insert_sources(
             if all(prev_div):
                 self.update_frame(self.source.repartition(divisions=prev_div))
             elif self.source.npartitions != prev_num:
-                self.source = self.source.repartition(npartitions=prev_num)
+                self.update_frame(self.source.repartition(npartitions=prev_num))
 
         return self
 

src/tape/ensemble_frame.py

@@ -666,7 +666,14 @@ def compute(self, **kwargs):
             self.ensemble._lazy_sync_tables_from_frame(self)
         return super().compute(**kwargs)
 
-    def repartition(self, **kwargs):
+    def repartition(
+        self,
+        divisions=None,
+        npartitions=None,
+        partition_size=None,
+        freq=None,
+        force=False,
+    ):
         """Repartition dataframe along new divisions
         
         Doc string below derived from dask.dataframe.DataFrame
@@ -725,7 +732,13 @@ def repartition(self, **kwargs):
         >>> df = df.repartition(divisions=[0, 5, 10, 20])  # doctest: +SKIP
         >>> df = df.repartition(freq='7d')  # doctest: +SKIP
         """
-        result = super().repartition(**kwargs)
+        result = super().repartition(
+            divisions=divisions, 
+            npartitions=npartitions, 
+            partition_size=partition_size,
+            freq=freq, 
+            force=force,
+        )
         return self._propagate_metadata(result)