考虑到前面question 中的代码，我认为最简单的解决方案是在for 循环内的AddFilenamesFn ParDo 内运行查询。请记住，beam.io.Read(beam.io.BigQuerySource(query=bqquery)) 用于将行作为源读取，而不是在中间步骤中。因此，在我建议的情况下，您可以直接使用 Python 客户端库 (google-cloud-bigquery>0.27.0)：
class AddFilenamesFn(beam.DoFn):
"""ParDo to output a dict with file id (retrieved from BigQuery) and row"""
def process(self, element, file_path):
from google.cloud import bigquery
client = bigquery.Client()
file_name = file_path.split("/")[-1]
query_job = client.query("""
SELECT FILE_ID
FROM test.file_mapping
WHERE FILENAME = '{0}'
LIMIT 1""".format(file_name))
results = query_job.result()
for row in results:
file_id = row.FILE_ID
yield {'filename':file_id, 'row':element}
这将是最直接的实施解决方案，但可能会出现问题。我们不是在管道开始时运行所有约 20 个可能的查询，而是为每行/记录运行一个查询。例如，如果我们在单个文件中有 3,000 个元素，则相同的查询将被启动 3,000 次。但是，每个不同的查询实际上应该只运行一次，随后的查询“重复”将命中cache。另请注意，缓存查询不会对交互式查询 limit 做出贡献。
我使用了之前answer的相同文件：
$ gsutil cat gs://$BUCKET/countries1.csv
id,country
1,sweden
2,spain
gsutil cat gs://$BUCKET/countries2.csv
id,country
3,italy
4,france
并添加了一个新表：
bq mk test.file_mapping FILENAME:STRING,FILE_ID:STRING
bq query --use_legacy_sql=false 'INSERT INTO test.file_mapping (FILENAME, FILE_ID) values ("countries1.csv", "COUNTRIES ONE"), ("countries2.csv", "COUNTRIES TWO")'
输出是：
INFO:root:{'filename': u'COUNTRIES ONE', 'row': u'id,country'}
INFO:root:{'filename': u'COUNTRIES ONE', 'row': u'1,sweden'}
INFO:root:{'filename': u'COUNTRIES ONE', 'row': u'2,spain'}
INFO:root:{'filename': u'COUNTRIES TWO', 'row': u'id,country'}
INFO:root:{'filename': u'COUNTRIES TWO', 'row': u'3,italy'}
INFO:root:{'filename': u'COUNTRIES TWO', 'row': u'4,france'}
另一种解决方案是使用beam.io.BigQuerySource() 加载所有表并将其具体化为侧面输入（取决于大小，这当然可能有问题），或者如您所说，将其分解为 N 个查询并保存每个查询到不同的侧输入。然后您可以为每条记录选择适当的记录并将其作为附加输入传递给AddFilenamesFn。尝试写那个也会很有趣。
我提出的第一个解决方案的完整代码：
import argparse, logging
from operator import add
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.io import ReadFromText
from apache_beam.io.filesystem import FileMetadata
from apache_beam.io.filesystem import FileSystem
from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
class GCSFileReader:
"""Helper class to read gcs files"""
def __init__(self, gcs):
self.gcs = gcs
class AddFilenamesFn(beam.DoFn):
"""ParDo to output a dict with file id (retrieved from BigQuery) and row"""
def process(self, element, file_path):
from google.cloud import bigquery
client = bigquery.Client()
file_name = file_path.split("/")[-1]
query_job = client.query("""
SELECT FILE_ID
FROM test.file_mapping
WHERE FILENAME = '{0}'
LIMIT 1""".format(file_name))
results = query_job.result()
for row in results:
file_id = row.FILE_ID
yield {'filename':file_id, 'row':element}
# just logging output to visualize results
def write_res(element):
logging.info(element)
return element
def run(argv=None):
parser = argparse.ArgumentParser()
known_args, pipeline_args = parser.parse_known_args(argv)
p = beam.Pipeline(options=PipelineOptions(pipeline_args))
gcs = GCSFileSystem(PipelineOptions(pipeline_args))
gcs_reader = GCSFileReader(gcs)
# in my case I am looking for files that start with 'countries'
BUCKET='BUCKET_NAME'
result = [m.metadata_list for m in gcs.match(['gs://{}/countries*'.format(BUCKET)])]
result = reduce(add, result)
# create each input PCollection name and unique step labels
variables = ['p{}'.format(i) for i in range(len(result))]
read_labels = ['Read file {}'.format(i) for i in range(len(result))]
add_filename_labels = ['Add filename {}'.format(i) for i in range(len(result))]
# load each input file into a separate PCollection and add filename to each row
for i in range(len(result)):
globals()[variables[i]] = p | read_labels[i] >> ReadFromText(result[i].path) | add_filename_labels[i] >> beam.ParDo(AddFilenamesFn(), result[i].path)
# flatten all PCollections into a single one
merged = [globals()[variables[i]] for i in range(len(result))] | 'Flatten PCollections' >> beam.Flatten() | 'Write results' >> beam.Map(write_res)
p.run()
if __name__ == '__main__':
run()