googlegenomics / gcp-variant-transforms / 920

pending completion

Build # 920

Build Type

push

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/googlegenomics/gcp-variant-transforms/commit/<a class=hub.com/googlegenomics/gcp-variant-transforms/commit/5197fb72ce3de2f67804bdbffb0fc2eb4707e5aa">5197fb72c<a href="https://github.com/googlegenomics/gcp-variant-transforms/commit/5197fb72ce3de2f67804bdbffb0fc2eb4707e5aa">">Merge </a><a class="double-link" href="https://github.com/googlegenomics/gcp-variant-transforms/commit/<a class="double-link" href="https://github.com/googlegenomics/gcp-variant-transforms/commit/f704b7332876ad7024300dd1c4ea6ecd8e90c5ab">f704b7332</a>">f704b7332</a><a href="https://github.com/googlegenomics/gcp-variant-transforms/commit/5197fb72ce3de2f67804bdbffb0fc2eb4707e5aa"> into aa0bf3115">aa0bf3115</a>

Run Details

125 of 125 new or added lines in 3 files covered. (100.0%)

5042 of 5552 relevant lines covered (90.81%)

0.91 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

56.04

/gcp_variant_transforms/transforms/variant_to_bigquery.py

# Copyright 2017 Google Inc.  All Rights Reserved.
#
# 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.

"""A PTransform to output a PCollection of ``Variant`` records to BigQuery."""

from __future__ import absolute_import

import exceptions
import random
import re
from typing import Dict, List  # pylint: disable=unused-import

import apache_beam as beam
from apache_beam.io.gcp.internal.clients import bigquery
from apitools.base.py import exceptions
from oauth2client.client import GoogleCredentials

from gcp_variant_transforms.beam_io import vcf_header_io  # pylint: disable=unused-import
from gcp_variant_transforms.libs import bigquery_row_generator
from gcp_variant_transforms.libs import bigquery_schema_descriptor  # pylint: disable=unused-import
from gcp_variant_transforms.libs import bigquery_util
from gcp_variant_transforms.libs import bigquery_vcf_schema
from gcp_variant_transforms.libs import processed_variant
from gcp_variant_transforms.libs import vcf_field_conflict_resolver
from gcp_variant_transforms.libs.variant_merge import variant_merge_strategy  # pylint: disable=unused-import
from gcp_variant_transforms.transforms import limit_write


# TODO(samanvp): remove this hack when BQ custom sink is added to Python SDK,
# see: https://issues.apache.org/jira/browse/BEAM-2801
# This has to be less than 10000.
_WRITE_SHARDS_LIMIT = 1000


@beam.typehints.with_input_types(processed_variant.ProcessedVariant)
class _ConvertToBigQueryTableRow(beam.DoFn):
  """Converts a ``Variant`` record to a BigQuery row."""

  def __init__(self,
               row_generator,
               allow_incompatible_records=False,
               omit_empty_sample_calls=False):
    # type: (bigquery_row_generator.BigQueryRowGenerator, bool, bool) -> None
    super(_ConvertToBigQueryTableRow, self).__init__()
    self._allow_incompatible_records = allow_incompatible_records
    self._omit_empty_sample_calls = omit_empty_sample_calls
    self._bigquery_row_generator = row_generator

  def process(self, record):
    return self._bigquery_row_generator.get_rows(
        record, self._allow_incompatible_records, self._omit_empty_sample_calls)


@beam.typehints.with_input_types(processed_variant.ProcessedVariant)
class VariantToBigQuery(beam.PTransform):
  """Writes PCollection of `ProcessedVariant` records to BigQuery."""

  def __init__(
      self,
      output_table,  # type: str
      header_fields,  # type: vcf_header_io.VcfHeader
      variant_merger=None,  # type: variant_merge_strategy.VariantMergeStrategy
      proc_var_factory=None,  # type: processed_variant.ProcessedVariantFactory
      append=False,  # type: bool
      update_schema=False,  # type: bool
      allow_incompatible_records=False,  # type: bool
      omit_empty_sample_calls=False,  # type: bool
      num_bigquery_write_shards=1  # type: int
      ):
    # type: (...) -> None
    """Initializes the transform.

    Args:
      output_table: Full path of the output BigQuery table.
      header_fields: Representative header fields for all variants. This is
        needed for dynamically generating the schema.
      variant_merger: The strategy used for merging variants (if any). Some
        strategies may change the schema, which is why this may be needed here.
      proc_var_factory: The factory class that knows how to convert Variant
        instances to ProcessedVariant. As a side effect it also knows how to
        modify BigQuery schema based on the ProcessedVariants that it generates.
        The latter functionality is what is needed here.
      append: If true, existing records in output_table will not be
        overwritten. New records will be appended to those that already exist.
      update_schema: If true, BigQuery schema will be updated by combining the
        existing schema and the new schema if they are compatible.
      allow_incompatible_records: If true, field values are casted to Bigquery
+       schema if there is a mismatch.
      omit_empty_sample_calls: If true, samples that don't have a given call
        will be omitted.
      num_bigquery_write_shards: If > 1, we will limit number of sources which
        are used for writing to the output BigQuery table.
    """
    self._output_table = output_table
    self._header_fields = header_fields
    self._variant_merger = variant_merger
    self._proc_var_factory = proc_var_factory
    self._append = append
    self._schema = bigquery_vcf_schema.generate_schema_from_header_fields(
        self._header_fields, self._proc_var_factory, self._variant_merger)
    # Resolver makes extra effort to resolve conflict when flag
    # allow_incompatible_records is set.
    self._bigquery_row_generator = bigquery_row_generator.BigQueryRowGenerator(
        bigquery_schema_descriptor.SchemaDescriptor(self._schema),
        vcf_field_conflict_resolver.FieldConflictResolver(
            resolve_always=allow_incompatible_records))

    self._allow_incompatible_records = allow_incompatible_records
    self._omit_empty_sample_calls = omit_empty_sample_calls
    self._num_bigquery_write_shards = num_bigquery_write_shards
    self._update_bigquery_schema(update_schema)

  def expand(self, pcoll):
    bq_rows = pcoll | 'ConvertToBigQueryTableRow' >> beam.ParDo(
        _ConvertToBigQueryTableRow(
            self._bigquery_row_generator,
            self._allow_incompatible_records,
            self._omit_empty_sample_calls))
    if self._num_bigquery_write_shards > 1:
      # We split data into self._num_bigquery_write_shards random partitions
      # and then write each part to final BQ by appending them together.
      # Combined with LimitWrite transform, this will avoid the BQ failure.
      bq_row_partitions = bq_rows | beam.Partition(
          lambda _, n: random.randint(0, n - 1),
          self._num_bigquery_write_shards)
      bq_writes = []
      for i in range(self._num_bigquery_write_shards):
        bq_rows = (bq_row_partitions[i] | 'LimitWrite' + str(i) >>
                   limit_write.LimitWrite(_WRITE_SHARDS_LIMIT))
        bq_writes.append(
            bq_rows | 'WriteToBigQuery' + str(i) >>
            beam.io.Write(beam.io.BigQuerySink(
                self._output_table,
                schema=self._schema,
                create_disposition=(
                    beam.io.BigQueryDisposition.CREATE_IF_NEEDED),
                write_disposition=(
                    beam.io.BigQueryDisposition.WRITE_APPEND))))
      return bq_writes
    else:
      return (bq_rows
              | 'WriteToBigQuery' >> beam.io.Write(beam.io.BigQuerySink(
                  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))))

  def _update_bigquery_schema(self, update_schema):
    # type: (bool) -> None
    if not update_schema or not self._schema:
      return
    # if table does not exist, do not need to update the schema.
    output_table_re_match = re.match(
        r'^((?P<project>.+):)(?P<dataset>\w+)\.(?P<table>[\w\$]+)$',
        self._output_table)
    credentials = GoogleCredentials.get_application_default().create_scoped(
        ['https://www.googleapis.com/auth/bigquery'])
    client = bigquery.BigqueryV2(credentials=credentials)
    try:
      project_id = output_table_re_match.group('project')
      dataset_id = output_table_re_match.group('dataset')
      table_id = output_table_re_match.group('table')
      found_table = client.tables.Get(bigquery.BigqueryTablesGetRequest(
          projectId=project_id,
          datasetId=dataset_id,
          tableId=table_id))
    except exceptions.HttpError:
      return

    new_schema = bigquery.TableSchema()
    new_schema.fields = _merge_field_schemas(self._schema.fields,
                                             found_table.schema.fields)
    original_schema = found_table.schema
    found_table.schema = new_schema
    try:
      client.tables.Update(bigquery.BigqueryTablesUpdateRequest(
          projectId=project_id,
          datasetId=dataset_id,
          table=found_table,
          tableId=table_id))
    except exceptions.HttpError:
      found_table.schema = original_schema


def _merge_field_schemas(
    field_schemas_1,  # type: List[bigquery.TableFieldSchema]
    field_schemas_2  # type: List[bigquery.TableFieldSchema]
    ):
  # type: (...) -> List[bigquery.TableFieldSchema]
  """Merges the `field_schemas_1` and `field_schemas_2`.

  Raises an error if there are fields with the same name, but different modes or
  different types.
  """
  existing_fields = {}  # type: Dict[str, bigquery.TableFieldSchema]
  merged_field_schemas = []  # type: List[bigquery.TableFieldSchema]
  for field_schema in field_schemas_1:
    existing_fields.update({field_schema.name: field_schema})
    merged_field_schemas.append(field_schema)

  for field_schema in field_schemas_2:
    if field_schema.name not in existing_fields.keys():
      merged_field_schemas.append(field_schema)
    else:
      existing_field_schema = existing_fields.get(field_schema.name)
      if field_schema.mode.lower() != existing_field_schema.mode.lower():
        raise ValueError(
            'The mode of field {} is not compatible. The original mode is {}, '
            'and the new mode is {}.'.format(field_schema.name,
                                             existing_field_schema.mode,
                                             field_schema.mode))
      if field_schema.type.lower() != existing_field_schema.type.lower():
        raise ValueError(
            'The type of field {} is not compatible. The original type is {}, '
            'and the new type is {}.'.format(field_schema.name,
                                             existing_field_schema.type,
                                             field_schema.type))
      if field_schema.type == bigquery_util.TableFieldConstants.TYPE_RECORD:
        existing_field_schema.fields = _merge_field_schemas(
            existing_field_schema.fields, field_schema.fields)
  return merged_field_schemas

1	# Copyright 2017 Google Inc. All Rights Reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	# http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	"""A PTransform to output a PCollection of ``Variant`` records to BigQuery."""	1✔
16
17	from __future__ import absolute_import	1✔
18
19	import exceptions	1✔
20	import random	1✔
21	import re	1✔
22	from typing import Dict, List # pylint: disable=unused-import	1✔
23
24	import apache_beam as beam	1✔
25	from apache_beam.io.gcp.internal.clients import bigquery	1✔
26	from apitools.base.py import exceptions	1✔
27	from oauth2client.client import GoogleCredentials	1✔
28
29	from gcp_variant_transforms.beam_io import vcf_header_io # pylint: disable=unused-import	1✔
30	from gcp_variant_transforms.libs import bigquery_row_generator	1✔
31	from gcp_variant_transforms.libs import bigquery_schema_descriptor # pylint: disable=unused-import	1✔
32	from gcp_variant_transforms.libs import bigquery_util	1✔
33	from gcp_variant_transforms.libs import bigquery_vcf_schema	1✔
34	from gcp_variant_transforms.libs import processed_variant	1✔
35	from gcp_variant_transforms.libs import vcf_field_conflict_resolver	1✔
36	from gcp_variant_transforms.libs.variant_merge import variant_merge_strategy # pylint: disable=unused-import	1✔
37	from gcp_variant_transforms.transforms import limit_write	1✔
38
39
40	# TODO(samanvp): remove this hack when BQ custom sink is added to Python SDK,
41	# see: https://issues.apache.org/jira/browse/BEAM-2801
42	# This has to be less than 10000.
43	_WRITE_SHARDS_LIMIT = 1000	1✔
44
45
46	@beam.typehints.with_input_types(processed_variant.ProcessedVariant)	1✔
47	class _ConvertToBigQueryTableRow(beam.DoFn):	1✔
48	"""Converts a ``Variant`` record to a BigQuery row."""
49
50	def __init__(self,	1✔
51	row_generator,
52	allow_incompatible_records=False,
53	omit_empty_sample_calls=False):
54	# type: (bigquery_row_generator.BigQueryRowGenerator, bool, bool) -> None
55	super(_ConvertToBigQueryTableRow, self).__init__()	1✔
56	self._allow_incompatible_records = allow_incompatible_records	1✔
57	self._omit_empty_sample_calls = omit_empty_sample_calls	1✔
58	self._bigquery_row_generator = row_generator	1✔
59
60	def process(self, record):	1✔
61	return self._bigquery_row_generator.get_rows(	1✔
62	record, self._allow_incompatible_records, self._omit_empty_sample_calls)
63
64
65	@beam.typehints.with_input_types(processed_variant.ProcessedVariant)	1✔
66	class VariantToBigQuery(beam.PTransform):	1✔
67	"""Writes PCollection of `ProcessedVariant` records to BigQuery."""
68
69	def __init__(	1✔
70	self,
71	output_table, # type: str
72	header_fields, # type: vcf_header_io.VcfHeader
73	variant_merger=None, # type: variant_merge_strategy.VariantMergeStrategy
74	proc_var_factory=None, # type: processed_variant.ProcessedVariantFactory
75	append=False, # type: bool
76	update_schema=False, # type: bool
77	allow_incompatible_records=False, # type: bool
78	omit_empty_sample_calls=False, # type: bool
79	num_bigquery_write_shards=1 # type: int
80	):
81	# type: (...) -> None
82	"""Initializes the transform.
83
84	Args:
85	output_table: Full path of the output BigQuery table.
86	header_fields: Representative header fields for all variants. This is
87	needed for dynamically generating the schema.
88	variant_merger: The strategy used for merging variants (if any). Some
89	strategies may change the schema, which is why this may be needed here.
90	proc_var_factory: The factory class that knows how to convert Variant
91	instances to ProcessedVariant. As a side effect it also knows how to
92	modify BigQuery schema based on the ProcessedVariants that it generates.
93	The latter functionality is what is needed here.
94	append: If true, existing records in output_table will not be
95	overwritten. New records will be appended to those that already exist.
96	update_schema: If true, BigQuery schema will be updated by combining the
97	existing schema and the new schema if they are compatible.
98	allow_incompatible_records: If true, field values are casted to Bigquery
99	+ schema if there is a mismatch.
100	omit_empty_sample_calls: If true, samples that don't have a given call
101	will be omitted.
102	num_bigquery_write_shards: If > 1, we will limit number of sources which
103	are used for writing to the output BigQuery table.
104	"""
105	self._output_table = output_table	×
106	self._header_fields = header_fields	×
107	self._variant_merger = variant_merger	×
108	self._proc_var_factory = proc_var_factory	×
109	self._append = append	×
110	self._schema = bigquery_vcf_schema.generate_schema_from_header_fields(	×
111	self._header_fields, self._proc_var_factory, self._variant_merger)
112	# Resolver makes extra effort to resolve conflict when flag
113	# allow_incompatible_records is set.
114	self._bigquery_row_generator = bigquery_row_generator.BigQueryRowGenerator(	×
115	bigquery_schema_descriptor.SchemaDescriptor(self._schema),
116	vcf_field_conflict_resolver.FieldConflictResolver(
117	resolve_always=allow_incompatible_records))
118
119	self._allow_incompatible_records = allow_incompatible_records	×
120	self._omit_empty_sample_calls = omit_empty_sample_calls	×
121	self._num_bigquery_write_shards = num_bigquery_write_shards	×
122	self._update_bigquery_schema(update_schema)	×
123
124	def expand(self, pcoll):	1✔
125	bq_rows = pcoll \| 'ConvertToBigQueryTableRow' >> beam.ParDo(	×
126	_ConvertToBigQueryTableRow(
127	self._bigquery_row_generator,
128	self._allow_incompatible_records,
129	self._omit_empty_sample_calls))
130	if self._num_bigquery_write_shards > 1:	×
131	# We split data into self._num_bigquery_write_shards random partitions
132	# and then write each part to final BQ by appending them together.
133	# Combined with LimitWrite transform, this will avoid the BQ failure.
134	bq_row_partitions = bq_rows \| beam.Partition(	×
135	lambda _, n: random.randint(0, n - 1),
136	self._num_bigquery_write_shards)
137	bq_writes = []	×
138	for i in range(self._num_bigquery_write_shards):	×
139	bq_rows = (bq_row_partitions[i] \| 'LimitWrite' + str(i) >>	×
140	limit_write.LimitWrite(_WRITE_SHARDS_LIMIT))
141	bq_writes.append(	×
142	bq_rows \| 'WriteToBigQuery' + str(i) >>
143	beam.io.Write(beam.io.BigQuerySink(
144	self._output_table,
145	schema=self._schema,
146	create_disposition=(
147	beam.io.BigQueryDisposition.CREATE_IF_NEEDED),
148	write_disposition=(
149	beam.io.BigQueryDisposition.WRITE_APPEND))))
150	return bq_writes	×
151	else:
152	return (bq_rows	×
153	\| 'WriteToBigQuery' >> beam.io.Write(beam.io.BigQuerySink(
154	self._output_table,
155	schema=self._schema,
156	create_disposition=(
157	beam.io.BigQueryDisposition.CREATE_IF_NEEDED),
158	write_disposition=(
159	beam.io.BigQueryDisposition.WRITE_APPEND
160	if self._append
161	else beam.io.BigQueryDisposition.WRITE_TRUNCATE))))
162
163	def _update_bigquery_schema(self, update_schema):	1✔
164	# type: (bool) -> None
165	if not update_schema or not self._schema:	×
166	return	×
167	# if table does not exist, do not need to update the schema.
168	output_table_re_match = re.match(	×
169	r'^((?P<project>.+):)(?P<dataset>\w+)\.(?P<table>[\w\$]+)$',
170	self._output_table)
171	credentials = GoogleCredentials.get_application_default().create_scoped(	×
172	['https://www.googleapis.com/auth/bigquery'])
173	client = bigquery.BigqueryV2(credentials=credentials)	×
174	try:	×
175	project_id = output_table_re_match.group('project')	×
176	dataset_id = output_table_re_match.group('dataset')	×
177	table_id = output_table_re_match.group('table')	×
178	found_table = client.tables.Get(bigquery.BigqueryTablesGetRequest(	×
179	projectId=project_id,
180	datasetId=dataset_id,
181	tableId=table_id))
182	except exceptions.HttpError:	×
183	return	×
184
185	new_schema = bigquery.TableSchema()	×
186	new_schema.fields = _merge_field_schemas(self._schema.fields,	×
187	found_table.schema.fields)
188	original_schema = found_table.schema	×
189	found_table.schema = new_schema	×
190	try:	×
191	client.tables.Update(bigquery.BigqueryTablesUpdateRequest(	×
192	projectId=project_id,
193	datasetId=dataset_id,
194	table=found_table,
195	tableId=table_id))
196	except exceptions.HttpError:	×
197	found_table.schema = original_schema	×
198
199
200	def _merge_field_schemas(	1✔
201	field_schemas_1, # type: List[bigquery.TableFieldSchema]
202	field_schemas_2 # type: List[bigquery.TableFieldSchema]
203	):
204	# type: (...) -> List[bigquery.TableFieldSchema]
205	"""Merges the `field_schemas_1` and `field_schemas_2`.
206
207	Raises an error if there are fields with the same name, but different modes or
208	different types.
209	"""
210	existing_fields = {} # type: Dict[str, bigquery.TableFieldSchema]	1✔
211	merged_field_schemas = [] # type: List[bigquery.TableFieldSchema]	1✔
212	for field_schema in field_schemas_1:	1✔
213	existing_fields.update({field_schema.name: field_schema})	1✔
214	merged_field_schemas.append(field_schema)	1✔
215
216	for field_schema in field_schemas_2:	1✔
217	if field_schema.name not in existing_fields.keys():	1✔
218	merged_field_schemas.append(field_schema)	1✔
219	else:
220	existing_field_schema = existing_fields.get(field_schema.name)	1✔
221	if field_schema.mode.lower() != existing_field_schema.mode.lower():	1✔
222	raise ValueError(	1✔
223	'The mode of field {} is not compatible. The original mode is {}, '
224	'and the new mode is {}.'.format(field_schema.name,
225	existing_field_schema.mode,
226	field_schema.mode))
227	if field_schema.type.lower() != existing_field_schema.type.lower():	1✔
228	raise ValueError(	1✔
229	'The type of field {} is not compatible. The original type is {}, '
230	'and the new type is {}.'.format(field_schema.name,
231	existing_field_schema.type,
232	field_schema.type))
233	if field_schema.type == bigquery_util.TableFieldConstants.TYPE_RECORD:	1✔
234	existing_field_schema.fields = _merge_field_schemas(	1✔
235	existing_field_schema.fields, field_schema.fields)
236	return merged_field_schemas	1✔

googlegenomics / gcp-variant-transforms / 920

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous