[arrow-datafusion-python] branch main updated: feature: Implement count method (#163)

[agrove]

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 ebf5ad6  feature: Implement count method (#163)
ebf5ad6 is described below

commit ebf5ad6c9a0e6fcbc976f07c205a283da07ba4e7
Author: Dejan Simic <10134699+sim...@users.noreply.github.com>
AuthorDate: Mon Feb 6 01:33:59 2023 +0100

    feature: Implement count method (#163)
---
 datafusion/tests/test_dataframe.py | 5 +++++
 src/dataframe.rs                   | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/datafusion/tests/test_dataframe.py 
b/datafusion/tests/test_dataframe.py
index a50532c..4d70845 100644
--- a/datafusion/tests/test_dataframe.py
+++ b/datafusion/tests/test_dataframe.py
@@ -453,3 +453,8 @@ def test_union_distinct(ctx):
 
 def test_cache(df):
     assert df.cache().collect() == df.collect()
+
+
+def test_count(df):
+    # Get number of rows
+    assert df.count() == 3
diff --git a/src/dataframe.rs b/src/dataframe.rs
index 3dd8210..2967a4f 100644
--- a/src/dataframe.rs
+++ b/src/dataframe.rs
@@ -312,4 +312,9 @@ impl PyDataFrame {
         wait_for_future(py, self.df.as_ref().clone().write_json(path))?;
         Ok(())
     }
+
+    // Executes this DataFrame to get the total number of rows.
+    fn count(&self, py: Python) -> PyResult<usize> {
+        Ok(wait_for_future(py, self.df.as_ref().clone().count())?)
+    }
 }