From 6a7c684d9b76e1dc4851096912fbb2ebf3c6a806 Mon Sep 17 00:00:00 2001
From: Ion Koutsouris <15728914+ion-elgreco@users.noreply.github.com>
Date: Thu, 25 Apr 2024 15:37:31 +0200
Subject: [PATCH] fix(python): reuse table state in write engine (#2453)

# Description
Instead of reusing the table state, it was being instantiated every time
you call write with the rust engine.

- https://github.com/delta-io/delta-rs/discussions/2448
---
 python/Cargo.toml              |  2 +-
 python/deltalake/_internal.pyi |  1 +
 python/deltalake/writer.py     |  1 +
 python/src/lib.rs              | 10 +++++++---
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/python/Cargo.toml b/python/Cargo.toml
index 1589ecb11e..0c205dc20f 100644
--- a/python/Cargo.toml
+++ b/python/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "deltalake-python"
-version = "0.17.1"
+version = "0.17.2"
 authors = ["Qingping Hou <dave2008713@gmail.com>", "Will Jones <willjones127@gmail.com>"]
 homepage = "https://github.com/delta-io/delta-rs"
 license = "Apache-2.0"
diff --git a/python/deltalake/_internal.pyi b/python/deltalake/_internal.pyi
index 7d2e6342a2..551528961e 100644
--- a/python/deltalake/_internal.pyi
+++ b/python/deltalake/_internal.pyi
@@ -174,6 +174,7 @@ def write_to_deltalake(
     data: pyarrow.RecordBatchReader,
     partition_by: Optional[List[str]],
     mode: str,
+    table: Optional[RawDeltaTable],
     schema_mode: Optional[str],
     predicate: Optional[str],
     name: Optional[str],
diff --git a/python/deltalake/writer.py b/python/deltalake/writer.py
index 2bc299215b..e4dd3fd42d 100644
--- a/python/deltalake/writer.py
+++ b/python/deltalake/writer.py
@@ -318,6 +318,7 @@ def write_deltalake(
             data=data,
             partition_by=partition_by,
             mode=mode,
+            table=table._table if table is not None else None,
             schema_mode=schema_mode,
             predicate=predicate,
             name=name,
diff --git a/python/src/lib.rs b/python/src/lib.rs
index 917bbb5750..5d3e813e33 100644
--- a/python/src/lib.rs
+++ b/python/src/lib.rs
@@ -1410,6 +1410,7 @@ fn write_to_deltalake(
     table_uri: String,
     data: PyArrowType<ArrowArrayStreamReader>,
     mode: String,
+    table: Option<&RawDeltaTable>,
     schema_mode: Option<String>,
     partition_by: Option<Vec<String>>,
     predicate: Option<String>,
@@ -1425,11 +1426,14 @@ fn write_to_deltalake(
         let save_mode = mode.parse().map_err(PythonError::from)?;
 
         let options = storage_options.clone().unwrap_or_default();
-        let table = rt()
-            .block_on(DeltaOps::try_from_uri_with_storage_options(
+        let table = if let Some(table) = table {
+            DeltaOps(table._table.clone())
+        } else {
+            rt().block_on(DeltaOps::try_from_uri_with_storage_options(
                 &table_uri, options,
             ))
-            .map_err(PythonError::from)?;
+            .map_err(PythonError::from)?
+        };
 
         let mut builder = table.write(batches).with_save_mode(save_mode);
         if let Some(schema_mode) = schema_mode {