From 57b4b7bfefeed8f14263bf69e3b73f423a69ba88 Mon Sep 17 00:00:00 2001
From: veenstrajelmer <60435591+veenstrajelmer@users.noreply.github.com>
Date: Wed, 14 Aug 2024 09:16:29 +0200
Subject: [PATCH] add unify_chunks() to merge_partitions()

---
 xugrid/ugrid/partitioning.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/xugrid/ugrid/partitioning.py b/xugrid/ugrid/partitioning.py
index 6b4797544..9af5d33f9 100644
--- a/xugrid/ugrid/partitioning.py
+++ b/xugrid/ugrid/partitioning.py
@@ -376,6 +376,9 @@ def merge_partitions(partitions, merge_ugrid_chunks: bool = True):
             )
             merged.update(merged_selection)
 
+    # prevent inconsistent chunks, for which merged.chunks will give a ValueError
+    merged = merged.unify_chunks()
+
     # Merge chunks along the UGRID dimensions.
     if merged.chunks and merge_ugrid_chunks:
         chunks = dict(merged.chunks)