Skip to content

Commit 5b96c04

Browse files
author
Oleksii Moskalenko
authored
Prevent race condition in BQ sink jobId generation (#877)
* fix race condition * [bq] temp file prefix in global window
1 parent 121cf20 commit 5b96c04

1 file changed

Lines changed: 55 additions & 68 deletions

File tree

‎storage/connectors/bigquery/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadsWithResult.java‎

Lines changed: 55 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,13 @@
55

66
import com.google.api.services.bigquery.model.TableRow;
77
import com.google.auto.value.AutoValue;
8-
98
import java.util.Collections;
109
import java.util.List;
1110
import java.util.Set;
1211
import javax.annotation.Nullable;
13-
12+
import org.apache.beam.sdk.Pipeline;
1413
import org.apache.beam.sdk.coders.*;
1514
import org.apache.beam.sdk.options.ValueProvider;
16-
import org.apache.beam.sdk.state.*;
1715
import org.apache.beam.sdk.transforms.*;
1816
import org.apache.beam.sdk.transforms.windowing.*;
1917
import org.apache.beam.sdk.values.*;
@@ -26,7 +24,7 @@
2624
@AutoValue
2725
public abstract class BatchLoadsWithResult<DestinationT>
2826
extends PTransform<
29-
PCollection<KV<DestinationT, TableRow>>, PCollection<KV<TableDestination, String>>> {
27+
PCollection<KV<DestinationT, TableRow>>, PCollection<KV<TableDestination, String>>> {
3028
static final Logger LOG = LoggerFactory.getLogger(BatchLoadsWithResult.class);
3129

3230
@VisibleForTesting
@@ -134,13 +132,14 @@ public PCollection<KV<TableDestination, String>> expand(
134132
input.getWindowingStrategy().getWindowFn() instanceof FixedWindows,
135133
"Input to BQ writer must be windowed in advance");
136134

137-
final PCollection<String> loadJobIdPrefixView = createLoadJobIdPrefixView(input);
135+
final PCollectionView<String> loadJobIdPrefixView = createLoadJobIdPrefixView(input);
138136
final PCollectionView<String> tempFilePrefixView =
139-
createTempFilePrefixView(loadJobIdPrefixView);
137+
createTempFilePrefixView(input.getPipeline());
140138

141139
PCollection<WriteBundlesToFiles.Result<DestinationT>> results =
142140
input
143-
.apply("WindowWithTrigger",
141+
.apply(
142+
"WindowWithTrigger",
144143
Window.<KV<DestinationT, TableRow>>configure()
145144
.triggering(
146145
Repeatedly.forever(
@@ -161,8 +160,8 @@ public void process(ProcessContext c) {
161160
.apply(
162161
"WriteGroupedRecords",
163162
ParDo.of(
164-
new WriteGroupedRecordsToFiles<>(
165-
tempFilePrefixView, DEFAULT_MAX_FILE_SIZE, getRowWriterFactory()))
163+
new WriteGroupedRecordsToFiles<>(
164+
tempFilePrefixView, DEFAULT_MAX_FILE_SIZE, getRowWriterFactory()))
166165
.withSideInputs(tempFilePrefixView))
167166
.setCoder(WriteBundlesToFiles.ResultCoder.of(getDestinationCoder()));
168167

@@ -176,8 +175,7 @@ tempFilePrefixView, DEFAULT_MAX_FILE_SIZE, getRowWriterFactory()))
176175
results
177176
.apply(
178177
Window.<WriteBundlesToFiles.Result<DestinationT>>configure()
179-
.triggering(DefaultTrigger.of())
180-
)
178+
.triggering(DefaultTrigger.of()))
181179
.apply("AttachSingletonKey", WithKeys.of((Void) null))
182180
.setCoder(
183181
KvCoder.of(
@@ -187,15 +185,15 @@ tempFilePrefixView, DEFAULT_MAX_FILE_SIZE, getRowWriterFactory()))
187185
.apply(
188186
"WritePartitionTriggered",
189187
ParDo.of(
190-
new WritePartition<>(
191-
false,
192-
getDynamicDestinations(),
193-
tempFilePrefixView,
194-
DEFAULT_MAX_FILES_PER_PARTITION,
195-
DEFAULT_MAX_BYTES_PER_PARTITION,
196-
multiPartitionsTag,
197-
singlePartitionTag,
198-
getRowWriterFactory()))
188+
new WritePartition<>(
189+
false,
190+
getDynamicDestinations(),
191+
tempFilePrefixView,
192+
DEFAULT_MAX_FILES_PER_PARTITION,
193+
DEFAULT_MAX_BYTES_PER_PARTITION,
194+
multiPartitionsTag,
195+
singlePartitionTag,
196+
getRowWriterFactory()))
199197
.withSideInputs(tempFilePrefixView)
200198
.withOutputTags(multiPartitionsTag, TupleTagList.of(singlePartitionTag)));
201199

@@ -206,82 +204,72 @@ tempFilePrefixView, DEFAULT_MAX_FILE_SIZE, getRowWriterFactory()))
206204
ShardedKeyCoder.of(NullableCoder.of(getDestinationCoder())),
207205
ListCoder.of(StringUtf8Coder.of())));
208206

209-
return writeSinglePartitionWithResult(
210-
partitions.get(singlePartitionTag), loadJobIdPrefixView.apply(View.asSingleton()));
207+
return writeSinglePartitionWithResult(partitions.get(singlePartitionTag), loadJobIdPrefixView);
211208
}
212209

213-
private PCollection<String> createLoadJobIdPrefixView(
210+
/**
211+
* Generates one jobId per window only if any feature row was submitted in this window. We need to
212+
* generate exactly one id per window, otherwise SingletonView will fail.
213+
*
214+
* @param input feature Rows
215+
* @return job id generated once per input's window
216+
*/
217+
private PCollectionView<String> createLoadJobIdPrefixView(
214218
PCollection<KV<DestinationT, TableRow>> input) {
215219
// We generate new JobId per each (input) window
216220
// To keep BQ job's name unique
217221
// Windowing of this generator is expected to be synchronized with input window
218222
// So generated ids can be applied as side input
223+
224+
String baseName = input.getPipeline().getOptions().getJobName().replaceAll("-", "");
225+
219226
return input
220227
.apply(
221-
"EraseKey",
228+
"EraseKeyAndValue",
222229
ParDo.of(
223-
new DoFn<KV<DestinationT, TableRow>, KV<Void, TableRow>>() {
230+
new DoFn<KV<DestinationT, TableRow>, String>() {
224231
@ProcessElement
225232
public void process(ProcessContext c) {
226-
c.output(KV.of(null, c.element().getValue()));
233+
// we don't need data, only fact of data existing
234+
c.output("");
227235
}
228236
}))
229237
.apply(
230-
"CreateJobId",
231-
ParDo.of(
232-
new DoFn<KV<Void, TableRow>, String>() {
233-
@StateId("oncePerWindow")
234-
private final StateSpec<SetState<Boolean>> oncePerWindow = StateSpecs.set(BooleanCoder.of());
235-
236-
@ProcessElement
237-
public void process(
238-
ProcessContext c,
239-
BoundedWindow w,
240-
@StateId("oncePerWindow") SetState<Boolean> oncePerWindow) {
241-
242-
// if set already contains something
243-
// it means we already generated Id for this window
244-
Boolean empty = oncePerWindow.isEmpty().read();
245-
if (empty != null && !empty) {
246-
return;
247-
}
248-
249-
// trying to add to Set and check if it was added
250-
// if true - we won and Id will be generated in current Process
251-
Boolean insertResult = oncePerWindow.addIfAbsent(true).read();
252-
if (insertResult != null && !insertResult) {
253-
return;
254-
}
255-
256-
c.output(
257-
String.format(
258-
"beam_load_%s_%s",
259-
c.getPipelineOptions().getJobName().replaceAll("-", ""),
260-
BigQueryHelpers.randomUUIDString()));
261-
262-
LOG.info("Pane {}, start: {}, last: {}", c.pane().getIndex(), c.pane().isFirst(), c.pane().isLast());
263-
LOG.info("[BQ] New window {}, {}", c.timestamp(), w.maxTimestamp());
264-
}
265-
}));
238+
Combine.globally(
239+
(SerializableFunction<Iterable<String>, String>)
240+
g ->
241+
String.format(
242+
"beam_load_%s_%s", baseName, BigQueryHelpers.randomUUIDString()))
243+
.withoutDefaults())
244+
.apply("JobIdView", View.asSingleton());
266245
}
267246

268-
private PCollectionView<String> createTempFilePrefixView(final PCollection<String> jobId) {
269-
return jobId
247+
/**
248+
* Generates one global (per all windows) prefix path to store files before load to BQ
249+
*
250+
* @param p Pipeline
251+
* @return view in global window
252+
*/
253+
private PCollectionView<String> createTempFilePrefixView(final Pipeline p) {
254+
return p.apply("CreateGlobalTempPrefix", Create.of(""))
270255
.apply(
271256
"GetTempFilePrefix",
272257
ParDo.of(
273258
new DoFn<String, String>() {
274259
@ProcessElement
275-
public void getTempFilePrefix(ProcessContext c, BoundedWindow w) {
260+
public void getTempFilePrefix(ProcessContext c) {
276261
String tempLocationRoot;
277262
if (getCustomGcsTempLocation() != null) {
278263
tempLocationRoot = getCustomGcsTempLocation().get();
279264
} else {
280265
tempLocationRoot = c.getPipelineOptions().getTempLocation();
281266
}
282267
String tempLocation =
283-
resolveTempLocation(tempLocationRoot, "BigQueryWriteTemp", c.element());
284-
LOG.info("[BQ] temp location generated {}, {}", tempLocation, w.maxTimestamp());
268+
resolveTempLocation(
269+
tempLocationRoot,
270+
"BigQueryWriteTemp",
271+
c.getPipelineOptions().getJobName());
272+
285273
c.output(tempLocation);
286274
}
287275
}))
@@ -321,5 +309,4 @@ PCollection<KV<TableDestination, String>> writeSinglePartitionWithResult(
321309
true,
322310
getSchemaUpdateOptions()));
323311
}
324-
325312
}

0 commit comments

Comments
 (0)