(wip) add retries to table write

Eventual-Inc · Jun 27, 2024 · db8fc53 · db8fc53
1 parent 7422f2f
commit db8fc53
Showing 1 changed file with 22 additions and 14 deletions.
diff --git a/daft/table/table_io.py b/daft/table/table_io.py
@@ -788,17 +788,25 @@ def _write_tabular_arrow_table(
     else:
         basename_template = f"{uuid4()}-{{i}}.{format.default_extname}"
 
-    pads.write_dataset(
-        arrow_table,
-        schema=schema,
-        base_dir=full_path,
-        basename_template=basename_template,
-        format=format,
-        partitioning=None,
-        file_options=opts,
-        file_visitor=file_visitor,
-        use_threads=True,
-        existing_data_behavior="overwrite_or_ignore",
-        filesystem=fs,
-        **kwargs,
-    )
+    for _ in range(3):
+        try:
+            pads.write_dataset(
+                arrow_table,
+                schema=schema,
+                base_dir=full_path,
+                basename_template=basename_template,
+                format=format,
+                partitioning=None,
+                file_options=opts,
+                file_visitor=file_visitor,
+                use_threads=True,
+                existing_data_behavior="overwrite_or_ignore",
+                filesystem=fs,
+                **kwargs,
+            )
+            break
+        except Exception as e:
+            print(f"!!!!! ERROR WRITING DATASET !!!!!: {e}")
+            error = e
+    else:
+        raise error