[python] write_dataset returns new dataset (#517)

lancedb · Feb 2, 2023 · 8274844 · 8274844
1 parent be10c4d
commit 8274844
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py
@@ -611,7 +611,7 @@ def write_dataset(
     mode: str = "create",
     max_rows_per_file: int = 1024 * 1024,
     max_rows_per_group: int = 1024,
-) -> bool:
+) -> LanceDataset:
     """
     Write a given data_obj to the given uri
 
@@ -652,4 +652,5 @@ def write_dataset(
     }
     if isinstance(uri, Path):
         uri = str(uri.absolute())
-    return _write_dataset(reader, str(uri), params)
+    _write_dataset(reader, str(uri), params)
+    return LanceDataset(str(uri))
diff --git a/python/python/tests/test_lance.py b/python/python/tests/test_lance.py
@@ -100,7 +100,5 @@ def l2sq(vec, mat):
 def test_count_rows(tmp_path):
     df = pd.DataFrame({"values": range(100)})
     tbl = pa.Table.from_pandas(df)
-    lance.write_dataset(tbl, tmp_path)
-
-    dataset = lance.dataset(tmp_path)
+    dataset = lance.write_dataset(tbl, tmp_path)
     assert dataset.count_rows() == 100