Fix rows limit bug on Dask

hi-primus · Oct 4, 2021 · 81b722a · 81b722a
1 parent 4709bda
commit 81b722a
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/optimus/engines/base/dask/rows.py b/optimus/engines/base/dask/rows.py
@@ -21,7 +21,11 @@ def limit(self, count):
         if count is None:
             return df
 
-        return self.root.new(self.root._base_to_dfd(df.cols.select("*").data.head(count), npartitions=-1))
+        return self.root.new(
+            self.root._base_to_dfd(
+                df.cols.select("*").data.head(count, npartitions=-1), df.partitions() if count > 100 else 1
+            )
+        )
 
     def _drop_duplicated_builtin(self, cols):