Analyse BigQuery data with TFDV (tensorflow data validation)

bq_tfdv.py

import apache_beam as beam
import pyarrow
import tensorflow_data_validation as tfdv
from tensorflow_metadata.proto.v0 import statistics_pb2
import numpy as np

pipeline_options = beam.pipeline.PipelineOptions.from_dictionary({
    'project': '[PROJECT_ID]'
})

def row_to_nparray(row: dict):
    return {k: np.asarray([v]) for k,v in row.items()}

with beam.Pipeline(options=pipeline_options) as p:
     r = (p 
            | 'BQ read' >> beam.io.Read(beam.io.BigQuerySource(query='SELECT * FROM `[TABLE]`',
                                                    use_standard_sql=True))
            | 'to dict' >> beam.Map(row_to_nparray)
            | ' batch' >> tfdv.utils.batch_util.BatchExamplesToArrowTables()
            | 'tdfv' >> tfdv.GenerateStatistics()
            | 'WriteStatsOutput' >> beam.io.WriteToTFRecord(
               'files', shard_name_template='',
               coder=beam.coders.ProtoCoder(
                   statistics_pb2.DatasetFeatureStatisticsList))
     )

     result = p.run()
     result.wait_until_finish()
     print(result)