Added chunking to CLI

Closes #69
Closes #118

Author: jeromekelleher

tests/test_cli.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2019 Tskit Developers
+# Copyright (c) 2019-2026 Tskit Developers
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -34,6 +34,7 @@
 
 import tszip
 import tszip.cli as cli
+from tszip import compat
 
 
 def get_stdout_for_pytest():
@@ -98,6 +99,7 @@ def test_default_values(self):
         self.assertEqual(args.decompress, False)
         self.assertEqual(args.list, False)
         self.assertEqual(args.stdout, False)
+        self.assertEqual(args.chunk_size, tszip.DEFAULT_CHUNK_SIZE)
         self.assertEqual(args.variants_only, False)
         self.assertEqual(args.suffix, ".tsz")
 
@@ -123,6 +125,14 @@ def test_decompress(self):
         args = parser.parse_args([infile, "--decompress"])
         self.assertTrue(args.decompress)
 
+    def test_chunk_size(self):
+        parser = cli.tszip_cli_parser()
+        infile = "tmp.trees.tsz"
+        args = parser.parse_args([infile, "-C", "1234"])
+        self.assertEqual(args.chunk_size, 1234)
+        args = parser.parse_args([infile, "--chunk-size=1234"])
+        self.assertTrue(args.chunk_size, 1234)
+
 
 class TestCli(unittest.TestCase):
     """
@@ -248,6 +258,20 @@ def test_variants_only(self):
         G2 = self.ts.genotype_matrix()
         self.assertTrue(np.array_equal(G1, G2))
 
+    def test_chunk_size(self):
+        self.assertTrue(self.trees_path.exists())
+        self.run_tszip([str(self.trees_path), "--chunk-size=20"])
+        self.assertFalse(self.trees_path.exists())
+        outpath = pathlib.Path(str(self.trees_path) + ".tsz")
+        self.assertTrue(outpath.exists())
+        ts = tszip.decompress(outpath)
+        self.assertEqual(ts.tables, self.ts.tables)
+        store = compat.create_zip_store(str(outpath), mode="r")
+        root = compat.create_zarr_group(store=store)
+        for _, g in root.groups():
+            for _, a in g.arrays():
+                assert a.chunks == (20,)
+
     def test_keep(self):
         self.assertTrue(self.trees_path.exists())
         self.run_tszip([str(self.trees_path), "--keep"])

tests/test_compression.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2021 Tskit Developers
+# Copyright (c) 2021-2026 Tskit Developers
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -307,7 +307,9 @@ def test_provenance(self):
                 root = compat.create_zarr_group(store=store)
                 self.assertEqual(
                     root.attrs["provenance"],
-                    provenance.get_provenance_dict({"variants_only": variants_only}),
+                    provenance.get_provenance_dict(
+                        {"variants_only": variants_only, "chunk_size": None}
+                    ),
                 )
 
     def write_file(self, attrs, path):
@@ -540,7 +542,7 @@ def test_good_chunks(self, tmpdir, chunk_size):
         ts2 = tszip.decompress(path)
         assert ts1 == ts2
 
-        store = compat.create_zip_store(path, mode="r")
+        store = compat.create_zip_store(str(path), mode="r")
         root = compat.create_zarr_group(store=store)
         for _, g in root.groups():
             for _, a in g.arrays():

tszip/__init__.py

@@ -21,6 +21,7 @@
 # SOFTWARE.
 from .compression import compress  # NOQA
 from .compression import decompress  # NOQA
+from .compression import DEFAULT_CHUNK_SIZE  # NOQA
 from .compression import load  # NOQA
 from .compression import print_summary  # NOQA
 from .provenance import __version__  # NOQA

tszip/cli.py

@@ -64,6 +64,14 @@ def tszip_cli_parser():
         "-v", "--verbosity", action="count", default=0, help="Increase the verbosity"
     )
     parser.add_argument("files", nargs="+", help="The files to compress/decompress.")
+    parser.add_argument(
+        "-C",
+        "--chunk-size",
+        type=int,
+        default=tszip.DEFAULT_CHUNK_SIZE,
+        help="Sets the size of array chunks to be compressed to the specified "
+        f"number of elements. Default={tszip.DEFAULT_CHUNK_SIZE}",
+    )
     parser.add_argument(
         "--variants-only",
         action="store_true",
@@ -125,7 +133,9 @@ def run_compress(args):
         check_output(outfile, args)
         if args.stdout:
             outfile = get_stdout()
-        tszip.compress(ts, outfile, variants_only=args.variants_only)
+        tszip.compress(
+            ts, outfile, variants_only=args.variants_only, chunk_size=args.chunk_size
+        )
         remove_input(infile, args)
 
 

tszip/compression.py

@@ -48,6 +48,8 @@
 FORMAT_NAME = "tszip"
 FORMAT_VERSION = [1, 0]
 
+DEFAULT_CHUNK_SIZE = 2**20
+
 
 def minimal_dtype(array):
     """
@@ -178,7 +180,7 @@ def compress_zarr(ts, root, variants_only=False, chunk_size=None):
     )
 
     if chunk_size is None:
-        chunk_size = 2**20
+        chunk_size = DEFAULT_CHUNK_SIZE
     if not isinstance(chunk_size, numbers.Integral):
         raise TypeError("Chunk size must be an integer")
     if chunk_size < 1: