数据流管道中的 TextIOWrapper 运行缓慢

Question

我正在将 csv 文件从 GCS 加载到 BigQuery 并通过 Cloud Composer 触发任务（然后接下来做一些其他事情）。由于某些字段中存在各种字符，bq load 命令无法正确解析文件，因此我转向 Dataflow 寻求帮助解析和加载。有 8 个文件，每个文件的大小约为 1GB。它有 96 列数据和约 300 万条记录，直接从 GCS 加载到 BQ。大多数字段都是 STRING，带有一些 NUMERIC 和 TIMESTAMP 类型。

我的管道运行，但非常缓慢。我可以成功地将文件读入 BigQuery，但管道会在 18 分钟后自动扩展到 +15 个工作人员，此时它只处理了大约 300k 行。 UI 显示它几乎没有推动 300 个元素/秒。

我已经尝试过在线发布的各种其他解决方案，但我需要未经编辑的数据（无法去除奇怪的字符），并且一些其他解决方案尝试使用 re 以逗号分隔，但是有在 STRING 字段中到处都是逗号，所以这对我不起作用。 （还有管道、制表符和任何潜在的字符，因此对其他内容进行定界也不是一个真正有用的选择）。该解决方案的优点在于能够使用apache_beam.Map 来并行化对记录执行的操作，但执行不正确会导致某些记录上的数据丢失或损坏，从而导致糟糕的结果。

csv 库是唯一能够始终如一地正确解析文件而不会丢失数据的库。所以我将打开的 GCS 文件传递​​给csv.DictReader，以便直接写入 BQ。无论我是在apache_beam.io.FileBasedSource 类中使用self.open_file() 方法还是在apache_beam.io.gcp.gcsio.GcsIO 类中使用open 方法，我都会得到一个_io.BufferedReader，它给我的是字节而不是字符串。所以我使用io.TextIOWrapper 来获取字符串而不是字节，这似乎“有效”但运行速度非常慢，如上所述。

我最初还尝试对 csv 文件进行 gzip 压缩并将其读入。我可以使用 gzip.open(_io.BufferedReader, 'rt') 而不是使用io.TextIOWrapper，这就像一个魅力。在这种情况下，管道始终运行并在大约 20 分钟内完成（根据数据流声称它可以做的事情，感觉仍然很长，但如果这是我能得到的最好的，那么我可以忍受）。 TextIOWrapper 似乎显着降低了它的速度（只是我的猜测），而 codecs 的其他解决方案似乎不起作用。

奇怪的是，即使使用TextIOWrapper，管道也可以在不到一分钟的时间内在 8 个 csv 文件之一上本地运行。所以现在我有点困惑。

（我也尝试在直接和肮脏的 python 中运行它，然后使用NLD_JSON 写入 bq load 并且工作但花了一个小时，这个过程由于各种原因不能超过。）

这是我的管道，其中注释部分显示了在 gzip 文件上执行所需的更改：

from __future__ import absolute_import

import argparse
from argparse import RawTextHelpFormatter
import logging

import apache_beam as beam
from apache_beam.io.gcp.bigquery import WriteToBigQuery
from apache_beam.io.filebasedsource import FileBasedSource

#############################################
# gzipped CSV Reading Class that converts to dictionary
#############################################


class MyCsvFileSource(FileBasedSource):
    def read_records(self, file_pattern, range_tracker):
        import os  # Need to import these inside of class otherwise the Pipeline will not recognize the library
        import csv
        from io import TextIOWrapper  # Comment this line out when reading gzipped csv
        #import gzip  # Uncomment this line when reading gzipped csv
        from apache_beam.io.gcp.gcsio import GcsIO
        my_gcs_io = GcsIO(storage_client=os.getenv('GOOGLE_APPLICATION_CREDENTIALS')) 
        # reader = csv.DictReader(gzip.open(my_gcs_io.open(filename=file_pattern, mode='r', mime_type='text/csv'), 'rt'))  # Uncomment this line when reading gzipped csv
        reader = csv.DictReader(TextIOWrapper(my_gcs_io.open(filename=file_pattern, mode='r', mime_type='text/csv')))  # Comment this line out when reading gzipped csv
        for record in reader:
            yield record

#############################################
# Define the Pipeline
#############################################


def run(argv=None):
    """This function defines the argument parser and pipeline arguments used to run the dataflow pipeline"""
    #############################################
    # Argument Parser
    #############################################

    parser = argparse.ArgumentParser(
        description="""
        This is an apache beam pipeline that will read a gzipped csv file and write to bigquery.
        The files can be read from GCS or local and written to bigquery in the same project.
        Required Pipeline Arguments:
        - runner
            To run locally specify the flag `--runner=DirectRunner`
            To run in GCP Dataflow specify the flag `--runner=DataflowRunner`
        - project [required only if accessing GCP, not required for local -> local]
            GCP Project ID where the Dataflow job will execute
            e.g. `--project=my-gcp-project`
        - stagingLocation [can specify local storage as well if running `DirectRunner`]
            Specify a GCS storage location where the Dataflow job can stage the code for workers to execute.
        - temp_location [can specify local storage as well if running `DirectRunner`]
            Specify a GCS storage location where the Dataflow job can stage the data for temporary storage.
        - subnetwork [required for reading from GCP GCS buckets]
            Need to specify a VPC subnetwork for the project using the following format
            `--subnetwork=regions/<REGION_NAME>/subnetworks/<SUB_NETWORK_NAME>`
        """,
        formatter_class=RawTextHelpFormatter)
    parser.add_argument("--input", help='The directory or filename that will be read into the pipeline containting 1 or more gzipped csv files')
    parser.add_argument("--output", help='The `dataset.table` where the records from `--input` will be written to')
    known_args, pipeline_args = parser.parse_known_args()

    #############################################
    # Dataflow Pipeline
    #############################################

    with beam.Pipeline(argv=pipeline_args) as p:
        (p
         | 'Read Files' >> beam.io.Read(MyCsvFileSource(known_args.input))
         | 'Write to BigQuery' >> WriteToBigQuery(table=known_args.output,
                                                  create_disposition=beam.io.BigQueryDisposition.CREATE_NEVER,
                                                  write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND))


if __name__ == '__main__':
    logging.getLogger().setLevel(logging.INFO)
    run()

我希望这条管道能够在 20 分钟内对从 GCS 到 BigQuery 的所有 8 个常规 csv 文件运行，并且在此过程中不会丢失任何记录。

希望大家能提供任何帮助。

Answer 1

** 不是您问题的答案，而是另一种方法 **

我了解到您正在尝试清理数据，同时加载到 BQ。您可能想探索 Cloud dataprep（在 GCP 控制台中的大数据部分下），它可以非常直观地清理您的数据和数据类型，例如（它是为转换您的数据而构建的）。然后，您可以将清理后的数据保存回 GCS，然后从 BQ UI 本身启动加载作业来填充您的 bigquery 表。