Support Snappy compression and a configurable GZIP level for make_examples output

deepvariant/BUILD

@@ -1564,6 +1564,7 @@ py_library(
     srcs = ["show_examples.py"],
     deps = [
         ":dv_constants",
+        ":dv_utils",
         "//third_party/nucleus/io:sharded_file_utils",
         "//third_party/nucleus/io:tfrecord",
         "//third_party/nucleus/protos:variants_py_pb2",

deepvariant/call_variants.py

@@ -511,12 +511,14 @@ def _parse_example_from_stream(blob):
         sharded_file_utils.normalize_to_sharded_file_pattern(path),
         shuffle=False,
     )
+    # All shards of one examples set share a codec; detect it from the path.
+    compression_type = dv_utils.compression_type_for_examples_path(path)
 
     def load_dataset(filename):
       dataset = tf.data.TFRecordDataset(
           filename,
           buffer_size=_DEFAULT_PREFETCH_BUFFER_BYTES,
-          compression_type='GZIP',
+          compression_type=compression_type,
       )
       return dataset
 

deepvariant/data_providers.py

@@ -216,12 +216,14 @@ def input_fn(
     file_list = tf.random.shuffle(file_list)
 
   ds = tf.data.Dataset.from_tensor_slices(file_list)
+  # All shards of one examples set share a codec; detect it from the path.
+  compression_type = dv_utils.compression_type_for_examples_path(path)
 
   def load_dataset(filename: str) -> tf.data.Dataset:
     return tf.data.TFRecordDataset(
         filename,
         buffer_size=config.prefetch_buffer_bytes,
-        compression_type='GZIP',
+        compression_type=compression_type,
     )
 
   ds = ds.interleave(

deepvariant/dv_utils.py

@@ -177,6 +177,24 @@ def example_sequencing_type(example):
   return example.features.feature['sequencing_type'].int64_list.value[0]
 
 
+def compression_type_for_examples_path(path: str) -> str:
+  """Returns the tf.data TFRecord compression_type for an examples path.
+
+  make_examples names Snappy-compressed output with a '.snappy' suffix; any
+  other suffix is treated as 'GZIP', matching the historical default in which
+  examples were always GZIP-compressed regardless of file name.
+
+  Args:
+    path: An examples file name, or a sharded/comma-separated pattern.
+
+  Returns:
+    'SNAPPY' if the (first) path ends (case-insensitively) in '.snappy',
+    otherwise 'GZIP'.
+  """
+  first = path.split(',')[0].strip().lower()
+  return 'SNAPPY' if first.endswith('.snappy') else 'GZIP'
+
+
 def get_one_example_from_examples_path(source, proto=None):
   """Get the first record from `source`.
 
@@ -195,7 +213,9 @@ def get_one_example_from_examples_path(source, proto=None):
         'Cannot find matching files with the pattern "{}"'.format(source)
     )
   dataset = tf.data.TFRecordDataset(
-      files, compression_type='GZIP', num_parallel_reads=tf.data.AUTOTUNE
+      files,
+      compression_type=compression_type_for_examples_path(files[0]),
+      num_parallel_reads=tf.data.AUTOTUNE,
   )
   if not proto:
     proto = example_pb2.Example

deepvariant/dv_utils_test.py

@@ -242,6 +242,63 @@ def test_no_cast(self):
     self.assertEqual(casted_images.dtype, tf.uint8)
 
 
+class CompressionTypeForExamplesPathTest(absltest.TestCase):
+
+  def test_snappy_suffix_selects_snappy(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path('examples.tfrecord.snappy'),
+        'SNAPPY',
+    )
+
+  def test_gz_suffix_selects_gzip(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path('examples.tfrecord.gz'),
+        'GZIP',
+    )
+
+  def test_no_compression_suffix_defaults_to_gzip(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path('examples.tfrecord'),
+        'GZIP',
+    )
+
+  def test_snappy_suffix_is_case_insensitive(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path('examples.tfrecord.SNAPPY'),
+        'SNAPPY',
+    )
+
+  def test_uses_first_of_comma_separated_paths(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path(
+            'shard0.tfrecord.snappy, shard1.tfrecord.snappy'
+        ),
+        'SNAPPY',
+    )
+
+  def test_sharded_snappy_pattern_selects_snappy(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path(
+            'examples-00000-of-00010.tfrecord.snappy'
+        ),
+        'SNAPPY',
+    )
+
+  def test_at_n_sharded_spec_selects_snappy(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path(
+            'examples@32.tfrecord.snappy'
+        ),
+        'SNAPPY',
+    )
+
+  def test_at_n_sharded_spec_gz_selects_gzip(self):
+    self.assertEqual(
+        dv_utils.compression_type_for_examples_path('examples@32.tfrecord.gz'),
+        'GZIP',
+    )
+
+
 if __name__ == '__main__':
   tf.config.run_functions_eagerly(True)
   absltest.main()

deepvariant/make_examples_core.py

@@ -1207,7 +1207,7 @@ def __init__(self, options, suffix=None):
 
     if options.examples_filename:
       clean_basename = re.sub(
-          r'(\@[0-9]+|-\*?\d*-of-\*?\d*|\.gz)',
+          r'(\@[0-9]+|-\*?\d*-of-\*?\d*|\.gz|\.snappy)',
           '',
           os.path.basename(options.examples_filename.lower()),
       )
@@ -1241,7 +1241,11 @@ def __init__(self, options, suffix=None):
       self._add_writer(
           'call_variant_outputs',
           dv_utils.get_tf_record_writer(
-              self._add_suffix(self.examples_filename, 'call_variant_outputs')
+              self._as_gzip_path(
+                  self._add_suffix(
+                      self.examples_filename, 'call_variant_outputs'
+                  )
+              )
           ),
       )
 
@@ -1289,7 +1293,9 @@ def __init__(self, options, suffix=None):
       self._add_writer(
           'small_model_examples',
           dv_utils.get_tf_record_writer(
-              self._add_suffix(self.examples_filename, 'small_model')
+              self._as_gzip_path(
+                  self._add_suffix(self.examples_filename, 'small_model')
+              )
           ),
       )
 
@@ -1309,6 +1315,23 @@ def _add_suffix(self, file_path, suffix):
     new_file = os.path.join(file_dir, new_file_base)
     return new_file
 
+  @staticmethod
+  def _as_gzip_path(file_path):
+    """Returns file_path with a '.gz' suffix in place of a '.snappy' one.
+
+    The auxiliary make_examples outputs (the small-model call_variant_outputs
+    and small_model_examples) are always GZIP-compressed, even when the main
+    examples output uses Snappy. They inherit their name from the examples
+    path, so swap a trailing '.snappy' for '.gz' to keep each file name
+    consistent with its actual codec. The suffix match is case-insensitive to
+    agree with the codec detection in
+    dv_utils.compression_type_for_examples_path and the C++ writer (both
+    lower-case the path before comparing).
+    """
+    if file_path.lower().endswith('.snappy'):
+      return file_path[: -len('.snappy')] + '.gz'
+    return file_path
+
   def write_examples(self, *examples):
     self._write('examples', *examples)
 

deepvariant/make_examples_core_test.py

@@ -367,6 +367,31 @@ def test_invalid_examples_filename_extension(self):
     ):
       make_examples_core.OutputsWriter(options)
 
+  def test_as_gzip_path_rewrites_snappy_suffix(self):
+    self.assertEqual(
+        make_examples_core.OutputsWriter._as_gzip_path(
+            'examples.tfrecord.snappy'
+        ),
+        'examples.tfrecord.gz',
+    )
+
+  def test_as_gzip_path_rewrites_uppercase_snappy_suffix(self):
+    # Codec detection is case-insensitive, so the side-output rename must be
+    # too; otherwise a '.SNAPPY' path yields GZIP bytes under a name readers
+    # treat as Snappy.
+    self.assertEqual(
+        make_examples_core.OutputsWriter._as_gzip_path(
+            'examples.tfrecord.SNAPPY'
+        ),
+        'examples.tfrecord.gz',
+    )
+
+  def test_as_gzip_path_passes_through_gz_suffix(self):
+    self.assertEqual(
+        make_examples_core.OutputsWriter._as_gzip_path('examples.tfrecord.gz'),
+        'examples.tfrecord.gz',
+    )
+
   @flagsaver.flagsaver
   def test_gvcf_output_enabled_is_false_without_gvcf_flag(self):
     FLAGS.mode = 'training'

deepvariant/make_examples_native.cc

@@ -148,8 +148,14 @@ ExamplesGenerator::ExamplesGenerator(
         LOG(INFO) << "Example filename not found for role: " << role;
         continue;
       }
-      sample.writer =
-          std::make_unique<nucleus::ExampleWriter>(it->second);
+      // An unset compression level (no explicit presence) maps to -1, which
+      // tells the writer to keep the library default rather than apply level 0.
+      const int compression_level =
+          options_.has_examples_compression_level()
+              ? options_.examples_compression_level()
+              : -1;
+      sample.writer = std::make_unique<nucleus::ExampleWriter>(
+          it->second, nucleus::ExampleFormat::kAuto, compression_level);
       if (!sample.writer->status().ok()) {
         LOG(FATAL) << "Failed to create Example writer for " << it->second;
       }

deepvariant/make_examples_options.py

@@ -92,6 +92,21 @@
     None,
     'Required. Path to write tf.Example protos in TFRecord format.',
 )
+_EXAMPLES_COMPRESSION_LEVEL = flags.DEFINE_integer(
+    'examples_compression_level',
+    -1,
+    (
+        'Compression level for the examples output, -1 or in [0, 9]. The'
+        ' examples codec is inferred from the --examples suffix (".snappy" ->'
+        ' Snappy, otherwise GZIP); this level only applies to GZIP output. A'
+        ' negative value (the default) uses the library default level.'
+    ),
+)
+flags.register_validator(
+    'examples_compression_level',
+    lambda v: v == -1 or 0 <= v <= 9,
+    message='--examples_compression_level must be -1 or in [0, 9].',
+)
 _CHECKPOINT = flags.DEFINE_string(
     'checkpoint',
     None,
@@ -1203,6 +1218,7 @@ def shared_flags_to_options(
         _OUTPUT_PHASING_ERROR_STATS.value or '',
     )
     options.examples_filename = examples
+    options.examples_compression_level = _EXAMPLES_COMPRESSION_LEVEL.value
     options.candidates_filename = candidates
     options.gvcf_filename = gvcf
     options.include_med_dp = _INCLUDE_MED_DP.value

deepvariant/protos/deepvariant.proto

@@ -780,6 +780,13 @@ message MakeExamplesOptions {
   string candidates_filename = 12;
   // Path to examples.
   string examples_filename = 13;
+  // Compression level for the examples output. Only applies to GZIP-compressed
+  // examples (i.e. a non-".snappy" examples_filename, whose codec is inferred
+  // from its suffix). A negative value uses the library default level (zlib
+  // Z_DEFAULT_COMPRESSION). Declared `optional` for explicit presence so that
+  // an unset field is distinguishable from a deliberate level 0 (store): when
+  // unset, the writer keeps the library default rather than disabling deflate.
+  optional int32 examples_compression_level = 102;
   // Path to a list of regions we are confident in, for determining which
   // candidate variants get labels.
   string confident_regions_filename = 14;

deepvariant/show_examples.py

@@ -53,6 +53,7 @@
 import tensorflow as tf
 
 from deepvariant import dv_constants
+from deepvariant import dv_utils
 from third_party.nucleus.io import sharded_file_utils
 from third_party.nucleus.io import tfrecord
 from third_party.nucleus.protos import variants_pb2
@@ -384,8 +385,15 @@ def run():
       tsv_df = pd.read_csv(_FILTER_BY_TSV.value, sep='\t', header=None)
       ids_from_tsv = set(tsv_df[0])
 
-    # Use nucleus.io.tfrecord to read all shards.
-    dataset = tfrecord.read_tfrecords(examples_path, compression_type='GZIP')
+    # Use nucleus.io.tfrecord to read all shards. The codec is inferred from the
+    # examples file-name suffix (".snappy" -> Snappy, otherwise GZIP), matching
+    # how make_examples named the output.
+    dataset = tfrecord.read_tfrecords(
+        examples_path,
+        compression_type=dv_utils.compression_type_for_examples_path(
+            examples_path
+        ),
+    )
 
     make_rgb = _IMAGE_TYPE.value in ['both', 'RGB']
     make_channels = _IMAGE_TYPE.value in ['both', 'channels']