File tree Expand file tree Collapse file tree 2 files changed +23
-0
lines changed Expand file tree Collapse file tree 2 files changed +23
-0
lines changed Original file line number Diff line number Diff line change @@ -236,3 +236,11 @@ def test_explain(df):
236
236
column ("a" ) - column ("b" ),
237
237
)
238
238
df .explain ()
239
+
240
+
241
+ def test_repartition (df ):
242
+ df .repartition (2 )
243
+
244
+
245
+ def test_repartition_by_hash (df ):
246
+ df .repartition_by_hash (column ("a" ), num = 2 )
Original file line number Diff line number Diff line change @@ -22,6 +22,7 @@ use datafusion::arrow::pyarrow::PyArrowConvert;
22
22
use datafusion:: arrow:: util:: pretty;
23
23
use datafusion:: dataframe:: DataFrame ;
24
24
use datafusion:: logical_plan:: JoinType ;
25
+ use datafusion:: prelude:: * ;
25
26
use pyo3:: exceptions:: PyTypeError ;
26
27
use pyo3:: prelude:: * ;
27
28
use pyo3:: types:: PyTuple ;
@@ -170,4 +171,18 @@ impl PyDataFrame {
170
171
let batches = wait_for_future ( py, df. collect ( ) ) ?;
171
172
Ok ( pretty:: print_batches ( & batches) ?)
172
173
}
174
+
175
+ /// Repartition a `DataFrame` based on a logical partitioning scheme.
176
+ fn repartition ( & self , num : usize ) -> PyResult < Self > {
177
+ let new_df = self . df . repartition ( Partitioning :: RoundRobinBatch ( num) ) ?;
178
+ Ok ( Self :: new ( new_df) )
179
+ }
180
+
181
+ /// Repartition a `DataFrame` based on a logical partitioning scheme.
182
+ #[ args( args = "*" , num) ]
183
+ fn repartition_by_hash ( & self , args : Vec < PyExpr > , num : usize ) -> PyResult < Self > {
184
+ let expr = args. into_iter ( ) . map ( |py_expr| py_expr. into ( ) ) . collect ( ) ;
185
+ let new_df = self . df . repartition ( Partitioning :: Hash ( expr, num) ) ?;
186
+ Ok ( Self :: new ( new_df) )
187
+ }
173
188
}
You can’t perform that action at this time.
0 commit comments