diff --git a/deepvariant/BUILD b/deepvariant/BUILD index 55fab5ba..6a0d5c17 100644 --- a/deepvariant/BUILD +++ b/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", diff --git a/deepvariant/call_variants.py b/deepvariant/call_variants.py index 767c00be..57f65a38 100644 --- a/deepvariant/call_variants.py +++ b/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 diff --git a/deepvariant/data_providers.py b/deepvariant/data_providers.py index 2501c214..0e5d44e9 100644 --- a/deepvariant/data_providers.py +++ b/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( diff --git a/deepvariant/dv_utils.py b/deepvariant/dv_utils.py index 7acb1a69..0f8c1dba 100644 --- a/deepvariant/dv_utils.py +++ b/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 diff --git a/deepvariant/dv_utils_test.py b/deepvariant/dv_utils_test.py index 5a444b1b..56a18f49 100644 --- a/deepvariant/dv_utils_test.py +++ b/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() diff --git a/deepvariant/make_examples_core.py b/deepvariant/make_examples_core.py index b2f9139f..bf56ae50 100644 --- a/deepvariant/make_examples_core.py +++ b/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) diff --git a/deepvariant/make_examples_core_test.py b/deepvariant/make_examples_core_test.py index 86505bae..71e16bc2 100644 --- a/deepvariant/make_examples_core_test.py +++ b/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' diff --git a/deepvariant/make_examples_native.cc b/deepvariant/make_examples_native.cc index 1054a6e7..27389f70 100644 --- a/deepvariant/make_examples_native.cc +++ b/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(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( + it->second, nucleus::ExampleFormat::kAuto, compression_level); if (!sample.writer->status().ok()) { LOG(FATAL) << "Failed to create Example writer for " << it->second; } diff --git a/deepvariant/make_examples_options.py b/deepvariant/make_examples_options.py index f84a2c75..3f32c56c 100644 --- a/deepvariant/make_examples_options.py +++ b/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 diff --git a/deepvariant/protos/deepvariant.proto b/deepvariant/protos/deepvariant.proto index 095db6a1..2edcbf12 100644 --- a/deepvariant/protos/deepvariant.proto +++ b/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; diff --git a/deepvariant/show_examples.py b/deepvariant/show_examples.py index d4bf1627..8f54ac5f 100644 --- a/deepvariant/show_examples.py +++ b/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'] diff --git a/scripts/run_deepvariant.py b/scripts/run_deepvariant.py index b6e9c07b..04875ddf 100644 --- a/scripts/run_deepvariant.py +++ b/scripts/run_deepvariant.py @@ -240,6 +240,28 @@ class ModelType(enum.Enum): ' has to be flag_name=true or flag_name=false.' ), ) +_EXAMPLES_COMPRESSION = flags.DEFINE_enum( + 'examples_compression', + 'GZIP', + ['GZIP', 'SNAPPY'], + ( + 'Codec for the make_examples tf.Example output. SNAPPY trades larger' + ' intermediate files for less make_examples CPU, which can help when' + ' output throughput is generously provisioned. The side outputs (gVCF,' + ' call_variant_outputs, small_model) are always GZIP regardless of this' + ' flag. This only affects the standard DeepVariant pipeline; the' + ' DeepTrio, DeepSomatic, and pangenome wrappers always use GZIP.' + ), +) +_EXAMPLES_COMPRESSION_LEVEL = flags.DEFINE_integer( + 'examples_compression_level', + None, + ( + 'GZIP compression level for the make_examples output (-1, or 0..9;' + ' -1 selects the zlib default). Only applies when --examples_compression' + ' is GZIP; ignored for SNAPPY.' + ), +) # Optional flag for postprocess variants _POSTPROCESS_CPUS = flags.DEFINE_integer( 'postprocess_cpus', @@ -704,6 +726,21 @@ def check_flags(): ' models.' ) + if _EXAMPLES_COMPRESSION_LEVEL.value is not None: + level = _EXAMPLES_COMPRESSION_LEVEL.value + if level != -1 and not 0 <= level <= 9: + raise RuntimeError( + '--examples_compression_level must be -1 or in [0, 9], got' + f' {level}.' + ) + if _EXAMPLES_COMPRESSION.value == 'SNAPPY': + logging.warning( + '--examples_compression_level is set to %d but' + ' --examples_compression is SNAPPY; the level only applies to GZIP' + ' output and will be ignored.', + level, + ) + def get_model_ckpt(model_type: str, customized_model: str) -> str: """Return the path to the model checkpoint based on the input args.""" @@ -725,9 +762,20 @@ def create_all_commands_and_logfiles(intermediate_results_dir): 'gvcf.tfrecord@{}.gz'.format(_NUM_SHARDS.value), ) + # The examples codec is inferred downstream from the file-name suffix + # (.snappy -> Snappy, else GZIP), so selecting the suffix here is the single + # control point: the same `examples` path feeds both make_examples (writer) + # and call_variants (reader), keeping them aligned. The side outputs below + # stay .gz: make_examples force-renames its auxiliary outputs to .gz (they go + # through the Python TFRecord writer, which cannot emit Snappy). + examples_suffix = ( + 'snappy' if _EXAMPLES_COMPRESSION.value == 'SNAPPY' else 'gz' + ) examples = os.path.join( intermediate_results_dir, - 'make_examples.tfrecord@{}.gz'.format(_NUM_SHARDS.value), + 'make_examples.tfrecord@{}.{}'.format( + _NUM_SHARDS.value, examples_suffix + ), ) small_model_cvo_records = os.path.join( intermediate_results_dir, @@ -771,6 +819,15 @@ def create_all_commands_and_logfiles(intermediate_results_dir): if tf.io.gfile.exists(potential_json): model_ckpt_json = potential_json + # The GZIP level only applies to GZIP output; check_flags() has already + # warned if a level was supplied alongside SNAPPY. + examples_compression_level = None + if ( + _EXAMPLES_COMPRESSION.value == 'GZIP' + and _EXAMPLES_COMPRESSION_LEVEL.value is not None + ): + examples_compression_level = _EXAMPLES_COMPRESSION_LEVEL.value + commands.append( make_examples_command( ref=_REF.value, @@ -787,6 +844,7 @@ def create_all_commands_and_logfiles(intermediate_results_dir): haploid_contigs=_HAPLOID_CONTIGS.value, par_regions_bed=_PAR_REGIONS.value, output_local_read_phasing=local_read_phasing_tsv_files, + examples_compression_level=examples_compression_level, ) ) diff --git a/third_party/nucleus/io/BUILD b/third_party/nucleus/io/BUILD index ec2eeb5d..ebf002fe 100644 --- a/third_party/nucleus/io/BUILD +++ b/third_party/nucleus/io/BUILD @@ -1335,5 +1335,6 @@ cc_test( deps = [ ":example_writer", "@com_google_googletest//:gtest_main", + "@org_tensorflow//tensorflow/core:lib", ], ) diff --git a/third_party/nucleus/io/example_writer.cc b/third_party/nucleus/io/example_writer.cc index 45548883..56e57b2d 100644 --- a/third_party/nucleus/io/example_writer.cc +++ b/third_party/nucleus/io/example_writer.cc @@ -39,6 +39,7 @@ #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/ascii.h" +#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_replace.h" @@ -49,13 +50,23 @@ namespace nucleus { +// Returns the TfRecord compression codec implied by the file-name suffix: +// "SNAPPY" for a ".snappy" extension, otherwise "GZIP". This mirrors the +// reader-side autodetection in dv_utils.compression_type_for_examples_path so a +// single source of truth (the file name) drives both ends. +std::string CompressionTypeForPath(absl::string_view path) { + return absl::EndsWith(absl::AsciiStrToLower(path), ".snappy") ? "SNAPPY" + : "GZIP"; +} + ExampleFormat AutodetectFormat(absl::string_view path, ExampleFormat format) { if (format != ExampleFormat::kAuto) return format; std::string path_str = absl::AsciiStrToLower(path); std::string extension; - // Replace shard strings and .gz extension. - RE2::GlobalReplace(&path_str, R"(\@[0-9]+|-\*?\d*-of-\*?\d*|\.gz)", ""); + // Replace shard strings and the compression extension (.gz or .snappy). + RE2::GlobalReplace(&path_str, + R"(\@[0-9]+|-\*?\d*-of-\*?\d*|\.gz|\.snappy)", ""); RE2::PartialMatch(path_str, R"(\.(bagz|tfrecords?)$)", // extension &extension); @@ -88,7 +99,7 @@ class ExampleWriter::Impl { class ExampleWriter::TfRecordImpl : public ExampleWriter::Impl { public: - explicit TfRecordImpl(absl::string_view path) { + TfRecordImpl(absl::string_view path, int compression_level) { UpdateStatus(tensorflow::Env::Default()->NewWritableFile(std::string(path), &tf_file_)); if (ABSL_PREDICT_FALSE(!status().ok())) { @@ -96,14 +107,28 @@ class ExampleWriter::TfRecordImpl : public ExampleWriter::Impl { return; } - const tensorflow::io::RecordWriterOptions& options = - tensorflow::io::RecordWriterOptions::CreateRecordWriterOptions( - "GZIP"); + // The codec is inferred from the file-name suffix (".snappy" -> SNAPPY, + // otherwise GZIP) so it always agrees with what the readers autodetect. + const std::string codec = CompressionTypeForPath(path); + tensorflow::io::RecordWriterOptions options = + tensorflow::io::RecordWriterOptions::CreateRecordWriterOptions(codec); + // The level only applies to the zlib-based codecs (GZIP/ZLIB); a negative + // value leaves the library default (Z_DEFAULT_COMPRESSION) in place. Snappy + // has no notion of a level, so a level set alongside a ".snappy" path is + // ignored -- warn rather than silently dropping it. + if (compression_level >= 0) { + if (codec == "SNAPPY") { + LOG(WARNING) << "Ignoring compression level " << compression_level + << " for Snappy output " << path + << " (Snappy does not support compression levels)."; + } else { + options.zlib_options.compression_level = compression_level; + } + } tf_ = std::make_unique( tf_file_.get(), options); - } bool Add(absl::string_view value, @@ -128,7 +153,8 @@ class ExampleWriter::TfRecordImpl : public ExampleWriter::Impl { ExampleWriter::ExampleWriter(absl::string_view path, - ExampleFormat format) { + ExampleFormat format, + int compression_level) { std::filesystem::path p = std::filesystem::path(path); if (!std::filesystem::is_directory(p.parent_path()) && p.parent_path() != "") { @@ -143,7 +169,7 @@ ExampleWriter::ExampleWriter(absl::string_view path, return; case ExampleFormat::kTfRecord: LOG(INFO) << "Writing output using TfRecord"; - impl_ = std::make_unique(path); + impl_ = std::make_unique(path, compression_level); break; case ExampleFormat::kBagz: break; diff --git a/third_party/nucleus/io/example_writer.h b/third_party/nucleus/io/example_writer.h index a753e4d0..94669984 100644 --- a/third_party/nucleus/io/example_writer.h +++ b/third_party/nucleus/io/example_writer.h @@ -33,6 +33,7 @@ #include #include +#include #include "absl/status/status.h" #include "absl/strings/string_view.h" @@ -47,11 +48,23 @@ enum class ExampleFormat { kBagz = 2, }; +// Returns the TfRecord compression codec implied by a file-name suffix: +// "SNAPPY" for a (case-insensitive) ".snappy" extension, otherwise "GZIP". +// Mirrors the reader-side detection in +// dv_utils.compression_type_for_examples_path so the file name is the single +// source of truth for both writing and reading. +std::string CompressionTypeForPath(absl::string_view path); + // Local writer for records, supports only a single file. class ExampleWriter { public: + // The TfRecord compression codec is inferred from the file-name suffix: a + // ".snappy" extension selects SNAPPY, otherwise GZIP. compression_level only + // applies to GZIP; a negative value uses the library default + // (Z_DEFAULT_COMPRESSION). explicit ExampleWriter(absl::string_view path, - ExampleFormat format = ExampleFormat::kAuto); + ExampleFormat format = ExampleFormat::kAuto, + int compression_level = -1); ~ExampleWriter(); bool Add(absl::string_view value, absl::string_view chrom = {}, diff --git a/third_party/nucleus/io/example_writer_test.cc b/third_party/nucleus/io/example_writer_test.cc index 0c0905ab..ca1f7598 100644 --- a/third_party/nucleus/io/example_writer_test.cc +++ b/third_party/nucleus/io/example_writer_test.cc @@ -29,9 +29,14 @@ * POSSIBILITY OF SUCH DAMAGE. */ #include "third_party/nucleus/io/example_writer.h" +#include #include +#include #include +#include +#include "tensorflow/core/lib/io/record_reader.h" +#include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/test.h" namespace nucleus { @@ -64,4 +69,56 @@ TEST(ExampleWriterTest, TFWriterTest) { EXPECT_TRUE(writer.Add(output_record)); } +TEST(ExampleWriterTest, CompressionTypeForPathDetectsCodec) { + EXPECT_EQ(CompressionTypeForPath("examples.tfrecord.snappy"), "SNAPPY"); + EXPECT_EQ(CompressionTypeForPath("examples.tfrecord.SNAPPY"), "SNAPPY"); + EXPECT_EQ(CompressionTypeForPath("examples-00000-of-00010.tfrecord.snappy"), + "SNAPPY"); + EXPECT_EQ(CompressionTypeForPath("examples.tfrecord.gz"), "GZIP"); + EXPECT_EQ(CompressionTypeForPath("examples.tfrecord"), "GZIP"); +} + +TEST(ExampleWriterTest, SnappyRoundTrip) { + // A ".snappy" path must produce a file a SNAPPY RecordReader can decode, + // proving the writer's suffix-inferred codec agrees with the readers. + const std::string path = "/tmp/dv_example_writer_snappy.tfrecord.snappy"; + const std::vector records = {"alpha", "beta", "gamma"}; + { + ExampleWriter writer(path); + for (const std::string& r : records) EXPECT_TRUE(writer.Add(r)); + EXPECT_TRUE(writer.Close()); + } + std::unique_ptr file; + ASSERT_TRUE( + tensorflow::Env::Default()->NewRandomAccessFile(path, &file).ok()); + tensorflow::io::RecordReaderOptions opts = + tensorflow::io::RecordReaderOptions::CreateRecordReaderOptions("SNAPPY"); + tensorflow::io::RecordReader reader(file.get(), opts); + tensorflow::uint64 offset = 0; + tensorflow::tstring value; + for (const std::string& expected : records) { + ASSERT_TRUE(reader.ReadRecord(&offset, &value).ok()); + EXPECT_EQ(std::string(value.data(), value.size()), expected); + } +} + +TEST(ExampleWriterTest, GzipCompressionLevelIsApplied) { + // A compressible payload, so level 0 (store) yields a strictly larger file + // than level 9 -- proving the level reaches the zlib options. + const std::string record(4096, 'A'); + auto write_at_level = [&record](const std::string& path, int level) { + ExampleWriter writer(path, ExampleFormat::kAuto, level); + for (int i = 0; i < 50; ++i) EXPECT_TRUE(writer.Add(record)); + EXPECT_TRUE(writer.Close()); + }; + const std::string path0 = "/tmp/dv_example_writer_l0.tfrecord.gz"; + const std::string path9 = "/tmp/dv_example_writer_l9.tfrecord.gz"; + write_at_level(path0, 0); + write_at_level(path9, 9); + tensorflow::uint64 size0 = 0, size9 = 0; + ASSERT_TRUE(tensorflow::Env::Default()->GetFileSize(path0, &size0).ok()); + ASSERT_TRUE(tensorflow::Env::Default()->GetFileSize(path9, &size9).ok()); + EXPECT_GT(size0, size9); +} + } // namespace nucleus