#316 quick fix

nicolay-r · Apr 7, 2022 · 1f73444 · 1f73444
1 parent a9a6f8a
commit 1f73444
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/arekit/common/data/storages/base.py b/arekit/common/data/storages/base.py
@@ -146,12 +146,20 @@ def free(self):
         del self._df
         gc.collect()
 
+    @staticmethod
+    def __iter_core(df):
+        for row_index, row in df.iterrows():
+            yield row_index, row
+
+    def iter_shuffled(self):
+        shuffled_df = self._df.sample(frac=1)
+        return self.__iter_core(shuffled_df)
+
     # endregion
 
     # region base methods
 
     def __iter__(self):
-        for row_index, row in self._df.iterrows():
-            yield row_index, row
+        return self.__iter_core(self._df)
 
     # endregion