From 950e05da817fe6a6f802d42073935d9548bcb3f0 Mon Sep 17 00:00:00 2001
From: Chang She <759245+changhiskhan@users.noreply.github.com>
Date: Tue, 9 Jan 2024 19:33:03 -0800
Subject: [PATCH] feat(python): add count_rows with filter option (#801)

Closes #795
---
 python/lancedb/table.py    | 13 ++++++++++++-
 python/tests/test_table.py | 11 +++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/python/lancedb/table.py b/python/lancedb/table.py
index 7db8a3d1..ae959c8d 100644
--- a/python/lancedb/table.py
+++ b/python/lancedb/table.py
@@ -646,8 +646,19 @@ class LanceTable(Table):
         self._dataset.restore()
         self._reset_dataset()
 
+    def count_rows(self, filter: Optional[str] = None) -> int:
+        """
+        Count the number of rows in the table.
+
+        Parameters
+        ----------
+        filter: str, optional
+            A SQL where clause to filter the rows to count.
+        """
+        return self._dataset.count_rows(filter)
+
     def __len__(self):
-        return self._dataset.count_rows()
+        return self.count_rows()
 
     def __repr__(self) -> str:
         return f"LanceTable({self.name})"
diff --git a/python/tests/test_table.py b/python/tests/test_table.py
index c6d948cf..4c769009 100644
--- a/python/tests/test_table.py
+++ b/python/tests/test_table.py
@@ -597,3 +597,14 @@ def test_compact_cleanup(db):
 
     with pytest.raises(Exception, match="Version 3 no longer exists"):
         table.checkout(3)
+
+
+def test_count_rows(db):
+    table = LanceTable.create(
+        db,
+        "my_table",
+        data=[{"text": "foo", "id": 0}, {"text": "bar", "id": 1}],
+    )
+    assert len(table) == 2
+    assert table.count_rows() == 2
+    assert table.count_rows(filter="text='bar'") == 1