[arrow-datafusion-python] branch main updated: Add partition_count property to ExecutionPlan. (#346)

[ag...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-python.git


The following commit(s) were added to refs/heads/main by this push:
     new af1a8eb  Add partition_count property to ExecutionPlan. (#346)
af1a8eb is described below

commit af1a8ebade3d537ce24f6f01a037a56177a917fa
Author: Kyle Brooks <84...@users.noreply.github.com>
AuthorDate: Tue Apr 25 19:54:40 2023 -0400

    Add partition_count property to ExecutionPlan. (#346)
---
 datafusion/tests/test_dataframe.py | 3 +++
 examples/substrait.py              | 2 +-
 src/physical_plan.rs               | 5 +++++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/datafusion/tests/test_dataframe.py b/datafusion/tests/test_dataframe.py
index cd78f3c..221b0cc 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 515311d..c579751 100644
--- a/examples/substrait.py
+++ b/examples/substrait.py
@@ -23,7 +23,7 @@ from datafusion import substrait as ss
 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 340d527..4c35f3e 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> {