并行化从 BigQuery 到 Python 中的本地 JSON 文件的导出答案

【问题标题】：Parallelize the export from BigQuery to local JSON file in Python并行化从 BigQuery 到 Python 中的本地 JSON 文件的导出
【发布时间】：2021-10-13 13:37:11
【问题描述】：

我们在 BigQuery 中有一个表，我们需要将其导出到本地以换行符分隔的 JSON 文件中。使用 BigQuery 的导出到 GCS 功能是有问题的，因为它将整数类型转换为字符串，请参阅 export bigquery data to cloud storage, the integer field changs to string format but float format stays as numeric format 和 How to preserve integer data type when exporting to JSON?，我自己尝试过，但整数丢失了。我们提出了以下解决方案，它维护整数类型，但非常慢：

当前工作代码

bq = bigquery.Client()
our_query = "select * from our_project.our_dataset.our_bq_table"

results_row_iter = bq.query(our_query) # google.bigquery.rowIterator
counter = 0
with open('/tmp/output_file.json', 'w') as f:
    for row in results_row_iter:
        f.write(json.dumps(dict(row), default=str) + '\n') # dumps as ndjson

our_bq_table 在 BigQuery 中为 5GB，有 340 万行和约 100 个字段，上面的 for 循环在我们的表上需要 90 分钟。 our_bq_table 在整数列 confId 上进行了分区，表中有大约 100 个唯一的 confId，值为 1 - 100。我们希望利用分区键 + 并行化来加速这个过程......不知何故。

我们要做什么的伪代码

bq = bigquery.Client()
base_query = "select * from our_project.our_dataset.our_bq_table"

all_conf_ids = range(1, 100)

def dump_conf_id(base_query, id):
    iter_query = f"{base_query} where confId = {id}"
    results_row_iter = bq.query(iter_query)
    counter = 0
    with open(f'output_file-{id}.json', 'w') as f:
        for row in results_row_iter:
            f.write(json.dumps(dict(row), default=str) + '\n') # dumps as ndjson

in parallel:
    for id in all_conf_ids:
        dump_conf_id(id)

# last step, perhaps concat the separate files into 1 somehow, assuming there are multiple output files...

这种方法利用了confId 字段，因此我们的 BigQuery 查询仍然很小。我不太确定如何在伪代码之外实现这一点，并且对弄清楚多线程、多处理和其他在 python 中并行化的方法感到不知所措。我们的最终输出需要是单个输出文件，伪代码转储到单独的文件中，但如果我们可以并行转储到单个文件中，那就太好了。

编辑：在实施解决方案之前我们试图解决的一个关键问题是我们应该为此使用多处理还是多线程，因为这 strong> 正在并行转储到本地 .json...

【问题讨论】：

也许您可以使用TO_JSON_STRING 函数导出到GCS，该函数保留类型。
@NickODell 这看起来非常有前途，我已经在 BigQuery 控制台中进行了测试，它似乎确实有效。但是，在 python 中运行bq.query(our_query) 之后，我很难将其保存到 JSON 文件中。我们试图不惜一切代价避免将这些结果转换为 pandas 数据框。
我们尝试过output = json.dumps(results_row_iter)，但收到错误TypeError: Object of type QueryJob is not JSON serializable

标签： python multithreading parallel-processing google-bigquery

【解决方案1】：

两条经验法则：

在 Python 中，您可以在程序很大程度上受 IO 限制的情况下使用多线程，但如果它受 CPU 限制，则必须使用多处理。这是因为 Python 的Global Interpreter Lock。
不应在同一个程序中混合使用多线程和多处理。

在这种情况下，我猜你的问题（为每个结果导出一个 JSON）是 CPU 密集型的，所以我建议使用多处理。这是一个代码示例：

from multiprocessing import Pool

n_jobs = 4
with Pool(n_jobs) as p:
    # call dump_conf_id on each member of all_conf_ids
    print(p.map(dump_conf_id, all_conf_ids))

[...] 但如果我们可以并行转储到单个文件中，那就太好了。

我不会费心尝试将它们写入单个文件。读取文件并连接它们可能是您正在做的最快的部分。一个快速的基准测试表明它以大约 780 MB/s 的速度运行：

import subprocess

with open("output.json", "wb") as f:
    filenames = [f'output_file-{id}.json' for i in all_conf_ids]
    subprocess.check_call(["cat", *filesnames], stdout=f)

【讨论】：