Set the zstd compression level when creating the VCF dataset (#568) (#569)

Co-authored-by: George Powley <george.powley@gmail.com>

Author: github-actions[bot]

apis/python/src/tiledbvcf/binding/libtiledbvcf.cc

@@ -114,5 +114,6 @@ PYBIND11_MODULE(libtiledbvcf, m) {
       .def("set_contig_mode", &Writer::set_contig_mode)
       .def("set_enable_allele_count", &Writer::set_enable_allele_count)
       .def("set_enable_variant_stats", &Writer::set_enable_variant_stats)
-      .def("set_compress_sample_dim", &Writer::set_compress_sample_dim);
+      .def("set_compress_sample_dim", &Writer::set_compress_sample_dim)
+      .def("set_compression_level", &Writer::set_compression_level);
 }

apis/python/src/tiledbvcf/binding/writer.cc

@@ -347,4 +347,9 @@ void Writer::set_compress_sample_dim(bool enable) {
       writer, tiledb_vcf_writer_set_compress_sample_dim(writer, enable));
 }
 
+void Writer::set_compression_level(int level) {
+  auto writer = ptr.get();
+  check_error(writer, tiledb_vcf_writer_set_compression_level(writer, level));
+}
+
 }  // namespace tiledbvcfpy

apis/python/src/tiledbvcf/binding/writer.h

@@ -235,6 +235,11 @@ class Writer {
   */
   void set_compress_sample_dim(bool enable);
 
+  /**
+    Set zstd compression level
+  */
+  void set_compression_level(int level);
+
  private:
   /** Helper function to free a C writer instance */
   static void deleter(tiledb_vcf_writer_t* w);

apis/python/src/tiledbvcf/dataset.py

@@ -569,6 +569,7 @@ def create_dataset(
         enable_allele_count: bool = True,
         enable_variant_stats: bool = True,
         compress_sample_dim: bool = True,
+        compression_level: int = 4,
     ):
         """
         Create a new dataset.
@@ -593,6 +594,8 @@ def create_dataset(
             Enable the variant stats ingestion task.
         compress_sample_dim
             Enable compression on the sample dimension.
+        compression_level
+            Compression level for zstd compression.
         """
         if self.mode != "w":
             raise Exception("Dataset not open in write mode")
@@ -629,6 +632,9 @@ def create_dataset(
         if compress_sample_dim is not None:
             self.writer.set_compress_sample_dim(compress_sample_dim)
 
+        if compression_level is not None:
+            self.writer.set_compression_level(compression_level)
+
         # This call throws an exception if the dataset already exists.
         self.writer.create_dataset()
 

apis/python/tests/test_tiledbvcf.py

@@ -1366,3 +1366,22 @@ def test_sample_compression(tmp_path, compress):
             found_zstd = found_zstd or "Zstd" in str(filter)
 
     assert found_zstd == compress
+
+
+@pytest.mark.parametrize("level", [1, 4, 16, 22])
+def test_compression_level(tmp_path, level):
+    # Create the dataset
+    dataset_uri = os.path.join(tmp_path, "compression_level")
+    array_uri = os.path.join(dataset_uri, "data")
+    ds = tiledbvcf.Dataset(dataset_uri, mode="w")
+    ds.create_dataset(compression_level=level)
+
+    check_if_compatible(array_uri)
+
+    # Check for the expected compression level
+    with tiledb.open(array_uri) as A:
+        for i in range(A.schema.nattr):
+            attr = A.schema.attr(i)
+            for filter in attr.filters:
+                if "Zstd" in str(filter):
+                    assert filter.level == level

libtiledbvcf/src/c_api/tiledbvcf.cc

@@ -1663,6 +1663,17 @@ int32_t tiledb_vcf_writer_set_compress_sample_dim(
   return TILEDB_VCF_OK;
 }
 
+int32_t tiledb_vcf_writer_set_compression_level(
+    tiledb_vcf_writer_t* writer, int level) {
+  if (sanity_check(writer) == TILEDB_VCF_ERR)
+    return TILEDB_VCF_ERR;
+
+  if (SAVE_ERROR_CATCH(writer, writer->writer_->set_compression_level(level)))
+    return TILEDB_VCF_ERR;
+
+  return TILEDB_VCF_OK;
+}
+
 /* ********************************* */
 /*               ERROR               */
 /* ********************************* */

libtiledbvcf/src/c_api/tiledbvcf.h

@@ -1602,6 +1602,14 @@ TILEDBVCF_EXPORT int32_t tiledb_vcf_writer_set_enable_variant_stats(
 TILEDBVCF_EXPORT int32_t tiledb_vcf_writer_set_compress_sample_dim(
     tiledb_vcf_writer_t* writer, bool enable);
 
+/**
+ * Sets zstd compression level
+ * @param writer VCF writer object
+ * @param level compression level
+ */
+TILEDBVCF_EXPORT int32_t
+tiledb_vcf_writer_set_compression_level(tiledb_vcf_writer_t* writer, int level);
+
 /* ********************************* */
 /*               ERROR               */
 /* ********************************* */

libtiledbvcf/src/cli/tiledbvcf.cc

@@ -441,6 +441,10 @@ void add_create(CLI::App& app) {
       args->compress_sample_dim,
       "Enable/disable compression of the sample dimension. Enabled by "
       "default.");
+  cmd->add_option(
+      "--compression-level",
+      args->compression_level,
+      "Set zstd compression level.");
 
   cmd->option_defaults()->group("Ingestion task options");
   cmd->add_flag(

libtiledbvcf/src/dataset/tiledbvcfdataset.cc

@@ -251,7 +251,8 @@ void TileDBVCFDataset::create(const CreationParams& params) {
       metadata,
       params.checksum,
       params.allow_duplicates,
-      params.compress_sample_dim);
+      params.compress_sample_dim,
+      params.compression_level);
 
   if (params.enable_allele_count) {
     AlleleCount::create(ctx, params.uri, params.checksum);
@@ -330,7 +331,8 @@ void TileDBVCFDataset::create_empty_data_array(
     const Metadata& metadata,
     const tiledb_filter_type_t& checksum,
     const bool allow_duplicates,
-    const bool compress_sample_dim) {
+    const bool compress_sample_dim,
+    const int compression_level) {
   ArraySchema schema(ctx, TILEDB_SPARSE);
   schema.set_capacity(metadata.tile_capacity);
   schema.set_order({{TILEDB_ROW_MAJOR, TILEDB_ROW_MAJOR}});
@@ -345,8 +347,14 @@ void TileDBVCFDataset::create_empty_data_array(
   FilterList float_attr_filters(ctx);
   FilterList byte_attr_filters(ctx);
 
+  // Use tile level filtering
+  str_attr_filters.set_max_chunk_size(0);
+  int_attr_filters.set_max_chunk_size(0);
+  float_attr_filters.set_max_chunk_size(0);
+  byte_attr_filters.set_max_chunk_size(0);
+
   Filter compression(ctx, TILEDB_FILTER_ZSTD);
-  compression.set_option(TILEDB_COMPRESSION_LEVEL, 4);
+  compression.set_option(TILEDB_COMPRESSION_LEVEL, compression_level);
 
   contig_coord_filters.add_filter({ctx, TILEDB_FILTER_RLE});
   pos_coord_filters.add_filter({ctx, TILEDB_FILTER_DOUBLE_DELTA})

libtiledbvcf/src/dataset/tiledbvcfdataset.h

@@ -66,6 +66,7 @@ struct CreationParams {
   bool enable_allele_count = true;
   bool enable_variant_stats = true;
   bool compress_sample_dim = true;
+  int compression_level = 4;
 };
 
 /** Arguments/params for dataset registration. */
@@ -872,7 +873,8 @@ class TileDBVCFDataset {
       const Metadata& metadata,
       const tiledb_filter_type_t& checksum,
       const bool allow_duplicates,
-      const bool compress_sample_dim);
+      const bool compress_sample_dim,
+      const int compression_level);
 
   /**
    * Creates the empty sample header array for a new dataset.