Merge pull request #122 from josephnowak/feature/improve-append-perfo…

…rmance2 Improve the performance of the append operation and also enable the s…
josephnowak · Sep 6, 2024 · 310a502 · 310a502
2 parents 52130e3 + 532e8c9
commit 310a502
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="TensorDB",
-    version="0.31.1",
+    version="0.31.2",
     description="Database based in a file system storage combined with Xarray and Zarr",
     author="Joseph Nowak",
     author_email="josephgonowak97@gmail.com",

diff --git a/tensordb/storages/zarr_storage.py b/tensordb/storages/zarr_storage.py
@@ -257,6 +257,7 @@ def append(
 
         rewrite = False
         act_coords = {k: coord for k, coord in act_data.indexes.items()}
+        data_to_append = {}
         slices_to_append = {}
         complete_data = act_data
 
@@ -279,9 +280,11 @@ def append(
                 k: slice(size, None) if k == dim else slice(0, size)
                 for k, size in complete_data.sizes.items()
             }
-            append_new_data = new_data.reindex(reindex_coords, fill_value=fill_value)
+            data_to_append[dim] = new_data.reindex(
+                reindex_coords, fill_value=fill_value
+            )
             complete_data = xr.concat(
-                [complete_data, append_new_data], dim=dim, fill_value=fill_value
+                [complete_data, data_to_append[dim]], dim=dim, fill_value=fill_value
             )
 
         complete_data = xr.Dataset(
@@ -305,17 +308,18 @@ def append(
             if dim not in slices_to_append:
                 continue
 
-            data_to_append = complete_data.isel(**slices_to_append[dim])
-            data_to_append = data_to_append.chunk(self.chunks)
-
+            data = data_to_append[dim].chunk(
+                complete_data.isel(**slices_to_append[dim]).chunks
+            )
             delayed_appends.append(
-                data_to_append.to_zarr(
+                data.to_zarr(
                     self.base_map,
                     append_dim=dim,
                     compute=compute,
                     synchronizer=self.synchronizer,
                     consolidated=True,
                     group=self.group,
+                    safe_chunks=False,
                 )
             )