diff --git a/datafusion/tests/test_dataframe.py b/datafusion/tests/test_dataframe.py
index cd78f3c88..221b0cc09 100644
--- a/datafusion/tests/test_dataframe.py
+++ b/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"
diff --git a/examples/substrait.py b/examples/substrait.py
index c167f7d90..3dcf10b21 100644
--- a/examples/substrait.py
+++ b/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"
 )
 
diff --git a/src/physical_plan.rs b/src/physical_plan.rs
index 340d527fa..4c35f3e60 100644
--- a/src/physical_plan.rs
+++ b/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> {