Add max-merge option to cooler load

open2c · Mar 23, 2024 · ded129e · ded129e
1 parent 9c9f864
commit ded129e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/cooler/cli/load.py b/src/cooler/cli/load.py
@@ -110,6 +110,13 @@
     "the platform default temp dir.",
     type=click.Path(exists=True, file_okay=False, dir_okay=True, allow_dash=True),
 )
+@click.option(
+    "--max-merge",
+    help="Maximum number of chunks to merge in a single pass.",
+    type=int,
+    default=200,
+    show_default=True,
+)
 @click.option(
     "--no-delete-temp",
     help="Do not delete temporary files when finished.",
@@ -146,6 +153,7 @@ def load(
     no_symmetric_upper,
     chunksize,
     mergebuf,
+    max_merge,
     temp_dir,
     no_delete_temp,
     storage_options,
@@ -343,9 +351,10 @@ def load(
         metadata=metadata,
         assembly=assembly,
         mergebuf=mergebuf,
-        ensure_sorted=False,
+        max_merge=max_merge,
         temp_dir=temp_dir,
         delete_temp=not no_delete_temp,
+        ensure_sorted=False,
         # boundscheck=True,
         # dupcheck=True,
         triucheck=True if symmetric_upper else False,

diff --git a/tests/test_create_ingest.py b/tests/test_create_ingest.py
@@ -415,6 +415,7 @@ def test_load_bg2_vs_coo():
         assembly="hg19",
         chunksize=int(20e6),
         mergebuf=int(15e6),
+        max_merge=200,
         field=(),
         count_as_float=False,
         one_based=False,
@@ -463,6 +464,7 @@ def test_load_zero_one_based_bg2():
         assembly="toy",
         chunksize=10,
         mergebuf=10,
+        max_merge=200,
         field=(),
         count_as_float=False,
         comment_char="#",
@@ -513,6 +515,7 @@ def test_load_zero_one_based_coo():
         assembly="toy",
         chunksize=10,
         mergebuf=10,
+        max_merge=200,
         field=(),
         count_as_float=False,
         comment_char="#",
@@ -616,6 +619,7 @@ def test_load_append_mode():
             assembly="hg19",
             chunksize=int(20e6),
             mergebuf=int(20e6),
+            max_merge=200,
             field=(),
             count_as_float=False,
             one_based=False,
@@ -656,6 +660,7 @@ def test_load_custom_tempdir():
             assembly="hg19",
             chunksize=int(20e6),
             mergebuf=int(20e6),
+            max_merge=200,
             field=(),
             count_as_float=False,
             one_based=False,