Scalable way to read large numbers of files with Apache Beam? - google-cloud-platform

Question

I’m writing a pipeline where I need to read the metadata files (500.000+ files) from the Sentinel2 dataset located on my Google Cloud Bucket with apache_beam.io.ReadFromTextWithFilename.
It works fine on a small subset, but when I ran it on the full dataset it seems to block on "Read Metadata" >> ReadFromTextWithFilename(f'gs://{BUCKET}/{DATA_FOLDER}/**/*metadata.json').
It dosen’t even show up in the Dataflow jobs list.
The pipeline looks like this:
with beam.Pipeline(options=pipeline_options) as pipeline:
meta = (
pipeline
| "Read Metadata" >> ReadFromTextWithFilename(f'gs://{BUCKET}/{DATA_FOLDER}/**/*metadata.json')
| "Extract metadata" >> beam.ParDo(ExtractMetaData())
)
table_spec = bigquery.TableReference(
datasetId="sentinel_2",
tableId="image_labels",
)
(
meta
| "Write To BigQuery" >> beam.io.WriteToBigQuery(
table_spec,
schema=table_schema(),
write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
)
)
I'm wondering:
Is there a smarter way to read large numbers of files?
Will copying the metadata files into one folder be more performant? (How much more does it cost to traverse sub-folders, as opposed to files in one folder)
Is the way to go to match the file names first with apache_beam.io.fileio.MatchAll and then read and extract in one or two following ParDos?

Answer 1

This is probably due to the pipeline running into Dataflow API limits when splitting the text source glob into a large number of sources.
Current solution is to use the transform ReadAllFromText which should not run into this.
In the future we hope to update transform ReadFromText for this case as well by using the Splittable DoFn framework.

Answer 2

It looks like I was suffering from a case of unwanted fusion.
Drawing inspiration from the page about file-processing on the Apache Beam website, I tried to add a Reshuffle to the pipeline.
I also upgraded to a paid Google Cloud account, thus getting higher quotas.
That resulted in Dataflow handling the job a lot better.
In fact Dataflow wanted to scale to 251 workers for my BiqQuery write job. At firste it didn't provision more workers, so I stopped the job and sat --num_workers=NUM_WORKERS and --max_num_workers=NUM_WORKERS, where NUM_WORKERS was the max qouta form my project. When running with those paramaters it scaled up automatically.
My final pipeline looks like this:
with beam.Pipeline(options=pipeline_options) as pipeline:
meta = (
pipeline
| MatchFiles(f'gs://{BUCKET}/{DATA_FOLDER}/*metadata.json')
| ReadMatches()
| beam.Reshuffle()
| beam.Map(lambda x: (x.metadata.path, x.read_utf8()))
| beam.ParDo(ExtractMetaData())
)
table_spec = bigquery.TableReference(
datasetId="sentinel_2",
tableId="image_labels",
)
(
meta
| "Write To BigQuery" >> beam.io.WriteToBigQuery(
table_spec,
schema=table_schema(),
write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE
)
)
Appendix
I also got a hint that SplitableParDos might be a solution, but I have not tested it.