add schema parameter to the repartition method (#192)

NVIDIA-Merlin · Dec 29, 2022 · 2fc6889 · 2fc6889
1 parent c2ef3f0
commit 2fc6889
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/merlin/io/dataset.py b/merlin/io/dataset.py
@@ -543,7 +543,8 @@ def repartition(self, npartitions=None, partition_size=None):
             .repartition(
                 npartitions=npartitions,
                 partition_size=partition_size,
-            )
+            ),
+            schema=self.schema,
         )
 
     @classmethod