From a88976e76dd5405e383e88633c278b6c55cb75de Mon Sep 17 00:00:00 2001 From: clarkzinzow Date: Mon, 25 Sep 2023 14:57:43 -0700 Subject: [PATCH] Fix num input partitions in coalesce. --- src/daft-plan/src/planner.rs | 2 +- tests/dataframe/test_repartition.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/daft-plan/src/planner.rs b/src/daft-plan/src/planner.rs index 5a38d0ed02..31b2f5e0fb 100644 --- a/src/daft-plan/src/planner.rs +++ b/src/daft-plan/src/planner.rs @@ -171,7 +171,7 @@ pub fn plan(logical_plan: &LogicalPlan) -> DaftResult { let input_physical = plan(input)?; Ok(PhysicalPlan::Coalesce(Coalesce::new( input_physical.into(), - logical_plan.partition_spec().num_partitions, + input.partition_spec().num_partitions, *num_to, ))) } diff --git a/tests/dataframe/test_repartition.py b/tests/dataframe/test_repartition.py index 92c96e5721..88d6281088 100644 --- a/tests/dataframe/test_repartition.py +++ b/tests/dataframe/test_repartition.py @@ -7,3 +7,9 @@ def test_into_partitions_some_empty() -> None: data = {"foo": [1, 2, 3]} df = daft.from_pydict(data).into_partitions(32).collect() assert df.to_pydict() == data + + +def test_into_partitions_coalesce() -> None: + data = {"foo": list(range(100))} + df = daft.from_pydict(data).into_partitions(20).into_partitions(1).collect() + assert df.to_pydict() == data