def filter_table_schema(self, include_fields=None): if include_fields is None: schema = self.table_schema else: schema = TableSchema() schema.fields = [field for field in self.table_schema.fields if field.name in include_fields] return schema
def table_schema(): table_schema = TableSchema() table_schema.fields = [ table_field('total_score', 'INTEGER'), table_field('user', 'STRING') ] return table_schema
def table_schema(): table_schema = TableSchema() table_schema.fields = [ table_field('team', 'STRING'), table_field('total_score', 'INTEGER'), table_field('window_start', 'TIMESTAMP') ] return table_schema
def table_schema(): """ The schema of the table where we'll be writing the output data """ table_schema = TableSchema() table_schema.fields = [ table_field('data', 'STRING') ] return table_schema
def table_schema(): """ The schema of the table where we'll be writing the output data """ table_schema = TableSchema() table_schema.fields = [ table_field('team', 'STRING'), table_field('total_score', 'INTEGER'), table_field('window_start', 'TIMESTAMP') ] return table_schema
def generate_schema(self): from apache_beam.io.gcp.internal.clients.bigquery import TableFieldSchema from apache_beam.io.gcp.internal.clients.bigquery import TableSchema json_fields = [ TableFieldSchema(name='country_code', type='STRING', mode='NULLABLE'), TableFieldSchema(name='country', type='JSON', mode='NULLABLE'), TableFieldSchema( name='stats', type='STRUCT', mode='NULLABLE', fields=[ TableFieldSchema( name="gdp_per_capita", type='JSON', mode='NULLABLE'), TableFieldSchema( name="co2_emissions", type='JSON', mode='NULLABLE'), ]), TableFieldSchema( name='cities', type='STRUCT', mode='REPEATED', fields=[ TableFieldSchema( name="city_name", type='STRING', mode='NULLABLE'), TableFieldSchema(name="city", type='JSON', mode='NULLABLE'), ]), TableFieldSchema(name='landmarks', type='JSON', mode='REPEATED'), ] schema = TableSchema(fields=json_fields) return schema
def _get_bq_schema(fields): bq_fields = [] for k, v in fields.items(): bq_fields.append( TableFieldSchema(name=k, type=v, description='Field %s' % k)) bq_fields.append( TableFieldSchema(name='_RAWTIMESTAMP', type='TIMESTAMP', description='Injected timestamp')) return TableSchema(fields=bq_fields)
def run(args=None): pipeline_options = PipelineOptions(args) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options.view_as(SetupOptions).save_main_session = True normalize_options = pipeline_options.view_as(NormalizeOptions) gcp_options = pipeline_options.view_as(GoogleCloudOptions) d1, d2 = parse_date_range(normalize_options.date_range) helper = QueryHelper(table=normalize_options.source_table, first_date_ts=d1, last_date_ts=d2) select_fields = ['mmsi', 'timestamp', 'seg_id', 'shipname', 'callsign', 'imo'] where_sql = 'shipname is not null or callsign is not null or imo is not null' if normalize_options.mmsi_quotient > 1: where_sql = "hash(mmsi) % {} = 0 and ({})".format(normalize_options.mmsi_quotient, where_sql) source_schema = helper.filter_table_schema(select_fields) source = BigQuerySource(query=helper.build_query(include_fields=select_fields, where_sql=where_sql)) dest_schema = TableSchema(fields=source_schema.fields) dest_schema.fields.append(TableFieldSchema(name=NORMALIZED_SHIPNAME, type='STRING')) dest_schema.fields.append(TableFieldSchema(name=NORMALIZED_CALLSIGN, type='STRING')) dest_schema.fields.append(TableFieldSchema(name=VALID_IMO, type='INTEGER')) pipeline = beam.Pipeline(options=pipeline_options) ( pipeline | "ReadSource" >> ReadAsJSONDict(source) | "ConvertTimestamp" >> beam.ParDo(ParseBeamBQStrTimestampDoFn()) | "AddTimestamp" >> beam.ParDo(TimestampedValueDoFn()) | "NormalizeNames" >> beam.ParDo(NormalizeNamesDoFn()) | "WriteDest" >> WriteToBigQueryDatePartitioned( temp_gcs_location=gcp_options.temp_location, table=normalize_options.dest_table, schema=dest_schema, write_disposition=BigQueryDisposition.WRITE_TRUNCATE) ) result = pipeline.run() success_states = set([PipelineState.DONE]) if normalize_options.wait: result.wait_until_finish() else: success_states.add(PipelineState.RUNNING) return 0 if result.state in success_states else 1