Add partition_count property to ExecutionPlan.

datafusion/tests/test_dataframe.py

@@ -370,6 +370,9 @@ def test_execution_plan(aggregate_df):
 
     assert expected == plan.display()
 
+    # Check the number of partitions is as expected.
+    assert type(plan.partition_count) is int
+
     expected = (
         "ProjectionExec: expr=[c1@0 as c1, SUM(test.c2)@1 as SUM(test.c2)]\n"
         "  Aggregate: groupBy=[[test.c1]], aggr=[[SUM(test.c2)]]\n"

examples/substrait.py

@@ -23,7 +23,7 @@
 ctx = SessionContext()
 
 # Register table with context
-ctx.register_parquet(
+ctx.register_csv(
     "aggregate_test_data", "./testing/data/csv/aggregate_test_100.csv"
 )
 

src/physical_plan.rs

@@ -53,6 +53,11 @@ impl PyExecutionPlan {
         let d = displayable(self.plan.as_ref());
         format!("{}", d.indent())
     }
+
+    #[getter]
+    pub fn partition_count(&self) -> usize {
+        self.plan.output_partitioning().partition_count()
+    }
 }
 
 impl From<PyExecutionPlan> for Arc<dyn ExecutionPlan> {