From 9965b4564d164d5a31fda48a36b65fa40d170c1d Mon Sep 17 00:00:00 2001
From: Lei Xu <lei@lancedb.com>
Date: Thu, 1 Jun 2023 15:58:45 -0700
Subject: [PATCH] [Python] Support drop table (#123)

Closes #86
---
 python/lancedb/db.py    | 13 +++++++++++++
 python/tests/test_db.py | 23 +++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/python/lancedb/db.py b/python/lancedb/db.py
index bbbeecdc..45bf70c1 100644
--- a/python/lancedb/db.py
+++ b/python/lancedb/db.py
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+import os
 
 import pyarrow as pa
 
@@ -112,3 +113,15 @@ class LanceDBConnection:
         A LanceTable object representing the table.
         """
         return LanceTable(self, name)
+
+    def drop_table(self, name: str):
+        """Drop a table from the database.
+
+        Parameters
+        ----------
+        name: str
+            The name of the table.
+        """
+        filesystem, path = pa.fs.FileSystem.from_uri(self.uri)
+        table_path = os.path.join(path, name + ".lance")
+        filesystem.delete_dir(table_path)
diff --git a/python/tests/test_db.py b/python/tests/test_db.py
index 975100f3..1ec39dd9 100644
--- a/python/tests/test_db.py
+++ b/python/tests/test_db.py
@@ -97,3 +97,26 @@ def test_create_mode(tmp_path):
     )
     tbl = db.create_table("test", data=new_data, mode="overwrite")
     assert tbl.to_pandas().item.tolist() == ["fizz", "buzz"]
+
+
+def test_delete_table(tmp_path):
+    db = lancedb.connect(tmp_path)
+    data = pd.DataFrame(
+        {
+            "vector": [[3.1, 4.1], [5.9, 26.5]],
+            "item": ["foo", "bar"],
+            "price": [10.0, 20.0],
+        }
+    )
+    db.create_table("test", data=data)
+
+    with pytest.raises(Exception):
+        db.create_table("test", data=data)
+
+    assert db.table_names() == ["test"]
+
+    db.drop_table("test")
+    assert db.table_names() == []
+
+    db.create_table("test", data=data)
+    assert db.table_names() == ["test"]
\ No newline at end of file