From 2fc6889cc5e9c7fbe497622710bea18c3344e824 Mon Sep 17 00:00:00 2001
From: Sara Rabhi <sara.rabhi@gmail.com>
Date: Wed, 21 Dec 2022 14:06:22 -0500
Subject: [PATCH] add schema parameter to the `repartition` method (#192)

---
 merlin/io/dataset.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/merlin/io/dataset.py b/merlin/io/dataset.py
index 5ee24e4d6..fc1f49d1b 100644
--- 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