From 56052ba8995c5d689c10c14d77b8a47da3e8f90b Mon Sep 17 00:00:00 2001
From: mauriceweber <maurice.weber@hotmail.com>
Date: Fri, 12 Jul 2024 13:41:32 +0200
Subject: [PATCH] add retry for write timeout errors

---
 daft/table/table_io.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/daft/table/table_io.py b/daft/table/table_io.py
index bdad3ebef8..e4a2365bc6 100644
--- a/daft/table/table_io.py
+++ b/daft/table/table_io.py
@@ -811,6 +811,7 @@ def _write_tabular_arrow_table(
     NUM_TRIES = 3
     JITTER_MS = 2_500
     MAX_BACKOFF_MS = 20_000
+    RETRY_ERRORS = ("InvalidPart", "curlCode: 28, Timeout was reached")
 
     for attempt in range(NUM_TRIES):
         try:
@@ -830,7 +831,7 @@ def _write_tabular_arrow_table(
             )
             break
         except OSError as e:
-            if "InvalidPart" not in str(e):
+            if all(err_str not in str(e) for err_str in RETRY_ERRORS):
                 raise
 
             if attempt == NUM_TRIES - 1: