From 0ad02d4243d1e207f8e54c7b763a7d84f8143573 Mon Sep 17 00:00:00 2001
From: ion-elgreco <ioncjk@gmail.com>
Date: Fri, 27 Oct 2023 10:27:34 +0200
Subject: [PATCH] add write support explicitly for pyarrow dataset

---
 python/deltalake/writer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/deltalake/writer.py b/python/deltalake/writer.py
index 42cd0aa541..9646e324f6 100644
--- 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