add write support explicitly for pyarrow dataset

delta-io · Oct 27, 2023 · 0ad02d4 · 0ad02d4
1 parent e5d2da3
commit 0ad02d4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/deltalake/writer.py b/python/deltalake/writer.py
@@ -65,6 +65,7 @@ def write_deltalake(
     table_or_uri: Union[str, Path, DeltaTable],
     data: Union[
         "pd.DataFrame",
+        ds.Dataset,
         pa.Table,
         pa.RecordBatch,
         Iterable[pa.RecordBatch],
@@ -331,6 +332,8 @@ def validate_batch(batch: pa.RecordBatch) -> pa.RecordBatch:
             batch_iter = [data]
         elif isinstance(data, pa.Table):
             batch_iter = data.to_batches()
+        elif isinstance(data, ds.Dataset):
+            batch_iter = data.to_batches()
         else:
             batch_iter = data