[Bug]: KafkaIO could fail with BigQueryIO.Write.withAutoSharding()

[nbali]

What happened?

BigQueryIO.Write.FILE_LOADS.withAutoSharding() uses GroupIntoBatches.withShardedKey(), which uses 'workerUuid' and 'threadId' as the sharding key. According to my understanding the problem is that the Kafka consumer read in KafkaIO.Read for a single partition most likely happens without any parallelism on the same worker on the same thread as it's being read in a FIFO manner due to the offset. This essentially means that .withShardedKey() has no effect whatsoever.

Although there is a 'FILE_TRIGGERING_BATCHING_DURATION' with '1s' duration, and a 'FILE_TRIGGERING_RECORD_COUNT' with '500000' count - and both triggers grouping, it still means if we are under 500k elements, and under 1s it will try to fire them at once. It is totally possible - with sufficiently high throughput, or 'outputWithTimestamp' that we have 500k elements in a single sec). This could result in OOME.

We should also have a size limit, not only time and count.

beam/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java

Lines 638 to 651 in a60105a

	Duration maxBufferingDuration =
	options.getMaxBufferingDurationMilliSec() > 0
	? Duration.millis(options.getMaxBufferingDurationMilliSec())
	: FILE_TRIGGERING_BATCHING_DURATION;
	// In contrast to fixed sharding with user trigger, here we use a global window with default
	// trigger and rely on GroupIntoBatches transform to group, batch and at the same time
	// parallelize properly. We also ensure that the files are written if a threshold number of
	// records are ready. Dynamic sharding is achieved via the withShardedKey() option provided by
	// GroupIntoBatches.
	return input
	.apply(
	GroupIntoBatches.<DestinationT, ElementT>ofSize(FILE_TRIGGERING_RECORD_COUNT)
	.withMaxBufferingDuration(maxBufferingDuration)
	.withShardedKey())

Issue Priority

Priority: 2

Issue Component

Component: io-java-gcp

[nbali]

.take-issue

[nbali]

Note to self: the same happens with https://github.com/apache/beam/blob/master/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java