Create call info table (#511)

Create call info table
- Use flag --generate_call_info_table to form the table.
- call_info table contains call_id, call_name and file_path.

Author: allieychen

gcp_variant_transforms/beam_io/vcf_header_io.py

@@ -18,7 +18,7 @@
 
 from collections import OrderedDict
 from functools import partial
-from typing import Dict, Iterable  # pylint: disable=unused-import
+from typing import Dict, Iterable, List  # pylint: disable=unused-import
 import vcf
 
 import apache_beam as beam
@@ -62,6 +62,7 @@ def __init__(self,
                alts=None,  # type: OrderedDict[str, vcf.parser._Alt]
                formats=None,  # type: OrderedDict[str, vcf.parser._Format]
                contigs=None,  # type: OrderedDict[str, vcf.parser._Contig]
+               samples=None,  # type: List[str]
                file_path=None  # type: str
               ):
     # type: (...) -> None
@@ -77,6 +78,7 @@ def __init__(self,
       alts: A dictionary mapping alt keys to vcf alt metadata values.
       formats: A dictionary mapping format keys to vcf format metadata values.
       contigs: A dictionary mapping contig keys to vcf contig metadata values.
+      samples: A list of sample names.
       file_path: The full file path of the vcf file.
     """
     # type: OrderedDict[str, OrderedDict]
@@ -85,6 +87,7 @@ def __init__(self,
     self.alts = self._values_asdict(alts or {})
     self.formats = self._values_asdict(formats or {})
     self.contigs = self._values_asdict(contigs or {})
+    self.samples = samples
     self.file_path = file_path
 
   def __eq__(self, other):
@@ -145,6 +148,7 @@ def read_records(
                     alts=vcf_reader.alts,
                     formats=vcf_reader.formats,
                     contigs=vcf_reader.contigs,
+                    samples=vcf_reader.samples,
                     file_path=file_path)
 
   def _read_headers(self, file_path):

gcp_variant_transforms/beam_io/vcf_header_io_test.py

@@ -42,6 +42,7 @@ def _get_header_from_reader(vcf_reader, file_path=None):
                    alts=vcf_reader.alts,
                    formats=vcf_reader.formats,
                    contigs=vcf_reader.contigs,
+                   samples=vcf_reader.samples,
                    file_path=file_path)
 
 

gcp_variant_transforms/libs/hashing_util.py

@@ -0,0 +1,31 @@
+# Copyright 2019 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generates hashing code for string."""
+
+import json
+import sys
+from typing import List  # pylint: disable=unused-import
+
+import farmhash
+
+
+def generate_unsigned_hash_code(strings, max_hash_value=sys.maxint):
+  # type: (List[str], int) -> int
+  """Generates a forever-fixed hash code for `strings`.
+
+  The hash code generated is in the range [0, max_hash_value). Note that the
+  hash code generated by farmhash.fingerprint64 is unsigned.
+  """
+  return farmhash.fingerprint64(json.dumps(strings)) % max_hash_value

gcp_variant_transforms/libs/hashing_util_test.py

@@ -0,0 +1,30 @@
+# Copyright 2019 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from gcp_variant_transforms.libs import hashing_util
+
+
+class HashingUtilTest(unittest.TestCase):
+
+  def test_generate_unsigned_hash_code(self):
+    hash_code = hashing_util.generate_unsigned_hash_code(['gs://bucket/blob',
+                                                          'sample 1'])
+    self.assertEqual(hash_code, 797596659968939265)
+
+    hash_code = hashing_util.generate_unsigned_hash_code(['gs://bucket/blob',
+                                                          'sample 1'],
+                                                         1000)
+    self.assertEqual(hash_code, 72)

gcp_variant_transforms/libs/sample_info_table_schema_generator.py

@@ -0,0 +1,51 @@
+# Copyright 2019 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generates sample_info table schema."""
+
+from apache_beam.io.gcp.internal.clients import bigquery
+
+from gcp_variant_transforms.libs import bigquery_util
+
+SAMPLE_ID = 'sample_id'
+SAMPLE_NAME = 'sample_name'
+FILE_PATH = 'file_path'
+TABLE_SUFFIX = 'sample_info'
+
+
+def compose_table_name(base_name, suffix):
+  # type: (str, List[str]) -> str
+  return '_'.join([base_name, suffix])
+
+
+def generate_schema():
+  # type: () -> bigquery.TableSchema
+  schema = bigquery.TableSchema()
+  schema.fields.append(bigquery.TableFieldSchema(
+      name=SAMPLE_ID,
+      type=bigquery_util.TableFieldConstants.TYPE_INTEGER,
+      mode=bigquery_util.TableFieldConstants.MODE_NULLABLE,
+      description='An Integer that uniquely identifies a sample.'))
+  schema.fields.append(bigquery.TableFieldSchema(
+      name=SAMPLE_NAME,
+      type=bigquery_util.TableFieldConstants.TYPE_STRING,
+      mode=bigquery_util.TableFieldConstants.MODE_NULLABLE,
+      description=('Name of the sample as we read it from the VCF file.')))
+  schema.fields.append(bigquery.TableFieldSchema(
+      name=FILE_PATH,
+      type=bigquery_util.TableFieldConstants.TYPE_STRING,
+      mode=bigquery_util.TableFieldConstants.MODE_NULLABLE,
+      description=('Full file path on GCS of the sample.')))
+
+  return schema

gcp_variant_transforms/libs/sample_info_table_schema_generator_test.py

@@ -0,0 +1,29 @@
+# Copyright 2019 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for `sample_info_table_schema_generator` module."""
+
+import unittest
+
+from gcp_variant_transforms.libs import sample_info_table_schema_generator
+
+
+class GenerateSampleInfoTableSchemaTest(unittest.TestCase):
+
+  def test_generate_sample_info_table_schema(self):
+    schema = sample_info_table_schema_generator.generate_schema()
+    expected_fields = [sample_info_table_schema_generator.SAMPLE_ID,
+                       sample_info_table_schema_generator.SAMPLE_NAME,
+                       sample_info_table_schema_generator.FILE_PATH]
+    self.assertEqual(expected_fields, [field.name for field in schema.fields])

gcp_variant_transforms/options/variant_transform_options.py

@@ -22,6 +22,7 @@
 from gcp_variant_transforms.beam_io import vcfio
 from gcp_variant_transforms.libs import bigquery_sanitizer
 from gcp_variant_transforms.libs import bigquery_util
+from gcp_variant_transforms.libs import sample_info_table_schema_generator
 
 
 class VariantTransformsOptions(object):
@@ -144,6 +145,16 @@ def add_arguments(self, parser):
     parser.add_argument('--output_table',
                         default='',
                         help='BigQuery table to store the results.')
+    parser.add_argument(
+        '--generate_sample_info_table',
+        type='bool', default=False, nargs='?', const=True,
+        help=('If set to True, a sample info table with the name '
+              'output_table_ + {} will be created. This table contains a '
+              'unique sample_id for each sample read from the VCF file. This '
+              'sample_id can be used to distinguish between sample names. '
+              '[EXPERIMENTAL]'
+             ).format(sample_info_table_schema_generator.TABLE_SUFFIX))
+
     parser.add_argument(
         '--split_alternate_allele_info_fields',
         type='bool', default=True, nargs='?', const=True,
@@ -207,6 +218,14 @@ def validate(self, parsed_args, client=None):
       if parsed_args.update_schema_on_append:
         raise ValueError('--update_schema_on_append requires --append to be '
                          'true.')
+      if parsed_args.generate_sample_info_table:
+        bigquery_util.raise_error_if_table_exists(
+            client,
+            project_id,
+            dataset_id,
+            '_'.join([table_id,
+                      sample_info_table_schema_generator.TABLE_SUFFIX]))
+
       bigquery_util.raise_error_if_table_exists(client,
                                                 project_id,
                                                 dataset_id,

gcp_variant_transforms/options/variant_transform_options_test.py

@@ -93,7 +93,8 @@ def _make_args(self, args):
 
   def test_valid_table_path(self):
     args = self._make_args(['--append',
-                            '--output_table', 'project:dataset.table'])
+                            '--output_table', 'project:dataset.table',
+                            '--generate_sample_info_table'])
     client = mock.Mock()
     client.datasets.Get.return_value = bigquery.Dataset(
         datasetReference=bigquery.DatasetReference(
@@ -104,7 +105,16 @@ def test_existing_table(self):
     args = self._make_args(['--append', 'False',
                             '--output_table', 'project:dataset.table'])
     client = mock.Mock()
-    self.assertRaises(ValueError, self._options.validate, args, client)
+    with self.assertRaisesRegexp(ValueError, 'project:dataset.table '):
+      self._options.validate(args, client)
+
+    args = self._make_args(['--append', 'False',
+                            '--output_table', 'project:dataset.table',
+                            '--generate_sample_info_table'])
+    client = mock.Mock()
+    with self.assertRaisesRegexp(ValueError,
+                                 'project:dataset.table_sample_info'):
+      self._options.validate(args, client)
 
   def test_no_project(self):
     args = self._make_args(['--output_table', 'dataset.table'])

gcp_variant_transforms/transforms/sample_info_to_bigquery.py

@@ -0,0 +1,70 @@
+# Copyright 2019 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Union  # pylint: disable=unused-import
+
+import apache_beam as beam
+
+from gcp_variant_transforms.beam_io import vcf_header_io  # pylint: disable=unused-import
+from gcp_variant_transforms.libs import sample_info_table_schema_generator
+from gcp_variant_transforms.libs import hashing_util
+
+
+class ConvertSampleInfoToRow(beam.DoFn):
+  """Extracts sample info from `VcfHeader` and converts it to a BigQuery row."""
+
+  def process(self, vcf_header):
+    # type: (vcf_header_io.VcfHeader) -> Dict[str, Union[int, str]]
+    for sample in vcf_header.samples:
+      sample_id = hashing_util.generate_unsigned_hash_code(
+          [vcf_header.file_path, sample], max_hash_value=pow(2, 63))
+      row = {
+          sample_info_table_schema_generator.SAMPLE_ID: sample_id,
+          sample_info_table_schema_generator.SAMPLE_NAME: sample,
+          sample_info_table_schema_generator.FILE_PATH: vcf_header.file_path
+      }
+      yield row
+
+
+class SampleInfoToBigQuery(beam.PTransform):
+  """Writes sample info to BigQuery."""
+
+  def __init__(self, output_table_prefix, append=False):
+    # type: (str, Dict[str, str], bool) -> None
+    """Initializes the transform.
+
+    Args:
+      output_table_prefix: The prefix of the output BigQuery table.
+      append: If true, existing records in output_table will not be
+        overwritten. New records will be appended to those that already exist.
+    """
+    self._output_table = sample_info_table_schema_generator.compose_table_name(
+        output_table_prefix, sample_info_table_schema_generator.TABLE_SUFFIX)
+    self._append = append
+    self._schema = sample_info_table_schema_generator.generate_schema()
+
+  def expand(self, pcoll):
+    return (pcoll
+            | 'ConvertSampleInfoToBigQueryTableRow' >> beam.ParDo(
+                ConvertSampleInfoToRow())
+            | 'WriteSampleInfoToBigQuery' >> beam.io.WriteToBigQuery(
+                self._output_table,
+                schema=self._schema,
+                create_disposition=(
+                    beam.io.BigQueryDisposition.CREATE_IF_NEEDED),
+                write_disposition=(
+                    beam.io.BigQueryDisposition.WRITE_APPEND
+                    if self._append
+                    else beam.io.BigQueryDisposition.WRITE_TRUNCATE),
+                method=beam.io.WriteToBigQuery.Method.FILE_LOADS))

gcp_variant_transforms/transforms/sample_info_to_bigquery_test.py

@@ -0,0 +1,61 @@
+# Copyright 2019 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for `sample_info_to_bigquery` module."""
+
+import unittest
+
+from apache_beam import transforms
+from apache_beam.testing import test_pipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+
+from gcp_variant_transforms.beam_io import vcf_header_io
+from gcp_variant_transforms.libs import sample_info_table_schema_generator
+from gcp_variant_transforms.transforms import sample_info_to_bigquery
+
+
+class ConvertSampleInfoToRowTest(unittest.TestCase):
+
+  def test_convert_sample_info_to_row(self):
+    vcf_header_1 = vcf_header_io.VcfHeader(samples=['Sample 1', 'Sample 2'],
+                                           file_path='file_1')
+    vcf_header_2 = vcf_header_io.VcfHeader(samples=['Sample 1', 'Sample 2'],
+                                           file_path='file_2')
+    file_path_to_file_hash = {'file_1': 'hash_1',
+                              'file_2': 'hash_2'}
+    expected_rows = [
+        {sample_info_table_schema_generator.SAMPLE_ID: 5961690698012655974,
+         sample_info_table_schema_generator.SAMPLE_NAME: 'Sample 1',
+         sample_info_table_schema_generator.FILE_PATH: 'file_1'},
+        {sample_info_table_schema_generator.SAMPLE_ID: 5854056809620188906,
+         sample_info_table_schema_generator.SAMPLE_NAME: 'Sample 2',
+         sample_info_table_schema_generator.FILE_PATH: 'file_1'},
+        {sample_info_table_schema_generator.SAMPLE_ID: 5259968798637352651,
+         sample_info_table_schema_generator.SAMPLE_NAME: 'Sample 1',
+         sample_info_table_schema_generator.FILE_PATH: 'file_2'},
+        {sample_info_table_schema_generator.SAMPLE_ID: 6253115674664185777,
+         sample_info_table_schema_generator.SAMPLE_NAME: 'Sample 2',
+         sample_info_table_schema_generator.FILE_PATH: 'file_2'}
+    ]
+    pipeline = test_pipeline.TestPipeline()
+    bigquery_rows = (
+        pipeline
+        | transforms.Create([vcf_header_1, vcf_header_2])
+        | 'ConvertToRow'
+        >> transforms.ParDo(sample_info_to_bigquery.ConvertSampleInfoToRow(
+            file_path_to_file_hash)))
+
+    assert_that(bigquery_rows, equal_to(expected_rows))
+    pipeline.run()