/*
* Copyright 2022 Google LLC.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions
* are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
* this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright
* notice, this list of conditions and the following disclaimer in the
* documentation and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from this
* software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
* ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
* LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*
*/
#include "third_party/nucleus/io/variant_reader.h"
#include <string>
#include <gmock/gmock-generated-matchers.h>
#include <gmock/gmock-matchers.h>
#include <gmock/gmock-more-matchers.h>
#include "tensorflow/core/platform/test.h"
#include "absl/strings/str_cat.h"
#include "third_party/nucleus/io/tfrecord_writer.h"
#include "third_party/nucleus/protos/variants.pb.h"
#include "third_party/nucleus/testing/protocol-buffer-matchers.h"
namespace {
using nucleus::EqualsProto;
using testing::Pointee;
nucleus::genomics::v1::Variant VariantProto(const std::string& ref_name,
int start, int end = 0) {
nucleus::genomics::v1::Variant v;
v.set_reference_name(ref_name);
v.set_start(start);
if (end != 0) {
v.set_end(end);
}
return v;
}
std::string VariantStr(const std::string& ref_name, int start, int end = 0) {
return VariantProto(ref_name, start, end).SerializeAsString();
}
TEST(IndexedReaderTest, EmptyShards) {
std::string path_a = absl::StrCat(getenv("TEST_TMPDIR"), "/", "a.gz");
std::string path_b = absl::StrCat(getenv("TEST_TMPDIR"), "/", "b.gz");
auto writer_a = nucleus::TFRecordWriter::New(path_a, "GZIP");
auto writer_b = nucleus::TFRecordWriter::New(path_b, "GZIP");
writer_a->Close();
writer_b->Close();
absl::flat_hash_map<std::string, uint32_t> contig_index_map = {{}};
auto reader =
nucleus::ShardedVariantReader::Open({path_a, path_b}, contig_index_map);
EXPECT_EQ(reader->GetAndReadNext().variant, nullptr);
}
TEST(IndexedReaderTest, ReadsRecordsSingleShard) {
std::string path_a = absl::StrCat(getenv("TEST_TMPDIR"), "/", "a.gz");
auto writer_a = nucleus::TFRecordWriter::New(path_a, "GZIP");
writer_a->WriteRecord(VariantStr("ref_a", 1));
writer_a->WriteRecord(VariantStr("ref_b", 1));
writer_a->WriteRecord(VariantStr("ref_c", 1));
writer_a->Close();
absl::flat_hash_map<std::string, uint32_t> contig_index_map = {
{"ref_a", 0}, {"ref_b", 1}, {"ref_c", 2}};
auto reader = nucleus::ShardedVariantReader::Open({path_a}, contig_index_map);
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 1))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_b", 1))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_c", 1))));
}
TEST(IndexedReaderTest, RecordsAreGloballyOrdered) {
std::string path_a = absl::StrCat(getenv("TEST_TMPDIR"), "/", "a.gz");
auto writer_a = nucleus::TFRecordWriter::New(path_a, "GZIP");
writer_a->WriteRecord(VariantStr("ref_a", 2));
writer_a->WriteRecord(VariantStr("ref_c", 100));
writer_a->WriteRecord(VariantStr("ref_d", 1));
writer_a->Close();
std::string path_b = absl::StrCat(getenv("TEST_TMPDIR"), "/", "b.gz");
auto writer_b = nucleus::TFRecordWriter::New(path_b, "GZIP");
writer_b->WriteRecord(VariantStr("ref_a", 1));
writer_b->WriteRecord(VariantStr("ref_b", 1));
writer_b->WriteRecord(VariantStr("ref_d", 2));
writer_b->WriteRecord(VariantStr("ref_d", 3));
writer_b->WriteRecord(VariantStr("ref_d", 4));
writer_b->Close();
std::string path_empty = absl::StrCat(getenv("TEST_TMPDIR"), "/", "empty.gz");
auto writer_c = nucleus::TFRecordWriter::New(path_empty, "GZIP");
writer_c->Close();
absl::flat_hash_map<std::string, uint32_t> contig_index_map = {
{"ref_a", 0}, {"ref_b", 1}, {"ref_c", 2}, {"ref_d", 3}};
auto reader = nucleus::ShardedVariantReader::Open(
{path_a, path_b, path_empty}, contig_index_map);
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 1))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 2))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_b", 1))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_c", 100))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_d", 1))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_d", 2))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_d", 3))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_d", 4))));
}
TEST(IndexedReaderTest, RecordsAreGloballyOrderedWithinContigs) {
std::string path_a = absl::StrCat(getenv("TEST_TMPDIR"), "/", "a.gz");
auto writer_a = nucleus::TFRecordWriter::New(path_a, "GZIP");
writer_a->WriteRecord(VariantStr("ref_a", 2, 3));
writer_a->WriteRecord(VariantStr("ref_a", 2, 5));
writer_a->WriteRecord(VariantStr("ref_a", 3, 4));
writer_a->WriteRecord(VariantStr("ref_b", 1, 1));
writer_a->Close();
std::string path_b = absl::StrCat(getenv("TEST_TMPDIR"), "/", "b.gz");
auto writer_b = nucleus::TFRecordWriter::New(path_b, "GZIP");
writer_b->WriteRecord(VariantStr("ref_a", 2, 2));
writer_b->WriteRecord(VariantStr("ref_a", 2, 6));
writer_b->WriteRecord(VariantStr("ref_a", 3, 1));
writer_b->WriteRecord(VariantStr("ref_b", 2, 1));
writer_b->Close();
absl::flat_hash_map<std::string, uint32_t> contig_index_map = {{"ref_a", 0},
{"ref_b", 1}};
auto reader =
nucleus::ShardedVariantReader::Open({path_a, path_b}, contig_index_map);
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 2, 2))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 2, 3))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 2, 5))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 2, 6))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 3, 1))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_a", 3, 4))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_b", 1, 1))));
EXPECT_THAT(reader->GetAndReadNext().variant,
Pointee(EqualsProto(VariantProto("ref_b", 2, 1))));
}
TEST(IndexedReaderTest, ReturnsContigIndex) {
std::string path_a = absl::StrCat(getenv("TEST_TMPDIR"), "/", "a");
auto writer_a = nucleus::TFRecordWriter::New(path_a, "");
writer_a->WriteRecord(VariantStr("ref_a", 1));
writer_a->WriteRecord(VariantStr("ref_b", 2));
writer_a->Close();
absl::flat_hash_map<std::string, uint32_t> contig_index_map = {{"ref_a", 1},
{"ref_b", 2}};
auto reader = nucleus::ShardedVariantReader::Open({path_a}, contig_index_map);
EXPECT_THAT(reader->GetAndReadNext().contig_map_index, 1);
EXPECT_THAT(reader->GetAndReadNext().contig_map_index, 2);
}
} // namespace
