[BigQueryIO] fetch updated schema for newly created Storage API stream writers #33231

ahmedabu98 · 2024-11-26T22:24:09Z

This problem happens whenever a StreamWriter is created after the table's schema is updated. The new StreamWriter ignores the previously updated schema and ends up using the base schema, dropping the new columns.
Such a case can happen when a pipeline decides to increase its parallelism (e.g. autosharding of threads or autoscaling of workers) after a schema update happens. The increased parallelism creates new stream writers that continue using the base schema, ignoring the updated schema.

This PR fixes this by fetching the stream's schema whenever we create a new append client. Existing tests are modified to cover this case, and new tests are added for schema update with dynamic destinations.

ahmedabu98 · 2024-12-18T00:01:29Z

R: @Abacn
R: @reuvenlax

github-actions · 2024-12-18T00:02:43Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

…api_schemaupdate_dynamicdest

reuvenlax · 2024-12-18T00:22:34Z

...e-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java

@@ -1419,7 +1421,11 @@ public WriteStream createWriteStream(String tableUrn, WriteStream.Type type)

    @Override
    public @Nullable WriteStream getWriteStream(String writeStream) {
-      return newWriteClient.getWriteStream(writeStream);
+      return newWriteClient.getWriteStream(


Let's add a boolean parameter to the method, so we only return schema if requested

FYI in our code base we always perform this call for the sole purpose of fetching the schema

Correct, but if that's the case we should replace it with a getSchema method. getWriteStream may be used in the future for other reasons, and if we always fetch the schema it may cause performance issues.

reuvenlax · 2024-12-18T00:28:50Z

...atform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java

@@ -531,6 +532,30 @@ public void process(
                      element.getKey().getKey(), dynamicDestinations, datasetService);
              tableSchema = converter.getTableSchema();
              descriptor = converter.getDescriptor(false);
+
+              if (autoUpdateSchema) {


I'm not sure this is the ideal place to put this. getAppendClientInfo is called whenever the static cache is populated, meaning that on any worker restart, range move, etc. we'll be forced to call this API again. However we have persistent state in this DoFn, so we know if it's a "new" key or not. Can we use that to gate calling this method instead?

We should always perform this check before creating a new StreamWriter , regardless of the reason for its creation. The only exception is if we already have an updated schema stored in state (see first if block above). If I'm following correctly, this method (getAppendClientInfo) will always create a new stream writer.

Also note that the updated schema is ignored when the StreamWriter object's creation time is later than the updated schema's.
i.e. it doesn't matter when the WriteStream itself was created

so this doesn't quite solve the following race condition:

schema S is updated to S' before creating stream writer, we detect it in this codepath and store it in state.

StreamWriter is destroyed (either because it is idle and evicted from the cache, or because the worker crashes).

We later recreate the StreamWriter. Because we have an updated schema in cache, we don't execute this codepath.

In this case, I think we'll completely miss the new schema. IIUC the best way to address this race would be for BigQuery to provide us with a version of StreamWriter that returns schemas with new fields - i.e. not basing it off of creation time.

Correct, this if block would only be executed if we have not detected S' yet.

But if S' is detected and stored in state, we would execute the first if block (line 520 here) and use that stored value. We still end up using the updated S', no?

Unless:

S is updated to S'

current writer recognizes it and stores S' in state

writer crashes for whatever reason

S' is updated to S'' (before new writer is created)

new writer is created, using the schema stored in state: S'

This new writer would have no chance to recognize S'' because it was created after that update.

This is a much narrower edge-case though. It's a small subset of the bigger problem we currently have, which involves essentially any new StreamWriter.

Are we okay with merging this solution and continue to pursue BigQuery to improve things on their side?

ahmedabu98 · 2024-12-30T14:39:44Z

@Abacn can you take a look? Reuven is OOO for some time

Abacn

Thanks, I saw the existing comments mainly concerned on additional API calls.

how often will the additional call will be made. If it's not significant I think this LGTM
does this PR introduces new API call if schema update feature isn't enabled (reading the change I believe no, but good to confirm)

...platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java

ahmedabu98 · 2024-12-30T17:40:05Z

how often will the additional call will be made. If it's not significant I think this LGTM

There's a call made for every new StreamWriter, which normally would be every time a new stream is created. In other words, it's roughly equal to the number of concurrent connections created at a given time. I checked in with BigQuery folks and they aren't worried about this quota.

does this PR introduces new API call if schema update feature isn't enabled

No this is strictly for when autoUpdateSchema is enabled

ahmedabu98 · 2025-01-02T13:08:04Z

The test passed on Dataflow V2 postcommits

Abacn

Thank you!

reuvenlax · 2025-01-02T17:18:50Z

CHANGES.md

@@ -79,6 +79,8 @@
 ## Bugfixes

 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* Fixed EventTimeTimer ordering in Prism. ([#32222](https://github.com/apache/beam/issues/32222)).


this doesn't seem to be in this PR

Hmm must be from a HEAD sync. will remove, thanks

reuvenlax · 2025-01-02T17:19:51Z

...e-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java

@@ -1419,7 +1421,11 @@ public WriteStream createWriteStream(String tableUrn, WriteStream.Type type)

    @Override
    public @Nullable WriteStream getWriteStream(String writeStream) {
-      return newWriteClient.getWriteStream(writeStream);
+      return newWriteClient.getWriteStream(


Correct, but if that's the case we should replace it with a getSchema method. getWriteStream may be used in the future for other reasons, and if we always fetch the schema it may cause performance issues.

reuvenlax · 2025-01-02T17:26:31Z

...atform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java

@@ -531,6 +532,30 @@ public void process(
                      element.getKey().getKey(), dynamicDestinations, datasetService);
              tableSchema = converter.getTableSchema();
              descriptor = converter.getDescriptor(false);
+
+              if (autoUpdateSchema) {


so this doesn't quite solve the following race condition:

schema S is updated to S' before creating stream writer, we detect it in this codepath and store it in state.

StreamWriter is destroyed (either because it is idle and evicted from the cache, or because the worker crashes).

We later recreate the StreamWriter. Because we have an updated schema in cache, we don't execute this codepath.

In this case, I think we'll completely miss the new schema. IIUC the best way to address this race would be for BigQuery to provide us with a version of StreamWriter that returns schemas with new fields - i.e. not basing it off of creation time.

add dynamic dest test

6bbbabf

github-actions bot added java io gcp labels Nov 26, 2024

Abacn mentioned this pull request Nov 27, 2024

[Bug]: Newly created Storage API write streams do not recognize the previously updated table schema #33238

Open

17 tasks

fix and add some tests

ad879e6

ahmedabu98 marked this pull request as ready for review December 18, 2024 00:00

ahmedabu98 changed the title ~~Storage API Schema Update with dynamic destinations~~ [BigQueryIO] fetch updated schema for newly created Storage API stream writers Dec 18, 2024

ahmedabu98 added this to the 2.62.0 Release milestone Dec 18, 2024

ahmedabu98 added 2 commits December 17, 2024 19:14

Merge branch 'master' of https://github.com/ahmedabu98/beam into stor…

cf8b5aa

…api_schemaupdate_dynamicdest

add to changes.md

7d455f9

reuvenlax reviewed Dec 18, 2024

View reviewed changes

ahmedabu98 and others added 2 commits December 18, 2024 12:05

fix whitespace

5159a61

Merge branch 'master' into storapi_schemaupdate_dynamicdest

21e0585

liferoad requested a review from Abacn December 30, 2024 14:46

Abacn reviewed Dec 30, 2024

View reviewed changes

...platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java Show resolved Hide resolved

trigger postcommits

1777269

github-actions bot added the build label Dec 30, 2024

Abacn approved these changes Jan 2, 2025

View reviewed changes

reuvenlax reviewed Jan 2, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[BigQueryIO] fetch updated schema for newly created Storage API stream writers #33231

[BigQueryIO] fetch updated schema for newly created Storage API stream writers #33231

ahmedabu98 commented Nov 26, 2024 •

edited

Loading

ahmedabu98 commented Dec 18, 2024

github-actions bot commented Dec 18, 2024

reuvenlax Dec 18, 2024

ahmedabu98 Dec 18, 2024

reuvenlax Jan 2, 2025

reuvenlax Dec 18, 2024

ahmedabu98 Dec 18, 2024

ahmedabu98 Dec 18, 2024

reuvenlax Jan 2, 2025

ahmedabu98 Jan 3, 2025 •

edited

Loading

ahmedabu98 Jan 3, 2025 •

edited

Loading

ahmedabu98 Jan 3, 2025 •

edited

Loading

ahmedabu98 commented Dec 30, 2024

Abacn left a comment

ahmedabu98 commented Dec 30, 2024

ahmedabu98 commented Jan 2, 2025

Abacn left a comment

reuvenlax Jan 2, 2025

ahmedabu98 Jan 3, 2025

reuvenlax Jan 2, 2025

reuvenlax Jan 2, 2025

[BigQueryIO] fetch updated schema for newly created Storage API stream writers #33231

Are you sure you want to change the base?

[BigQueryIO] fetch updated schema for newly created Storage API stream writers #33231

Conversation

ahmedabu98 commented Nov 26, 2024 • edited Loading

ahmedabu98 commented Dec 18, 2024

github-actions bot commented Dec 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ahmedabu98 Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

ahmedabu98 Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

ahmedabu98 Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

ahmedabu98 commented Dec 30, 2024

Abacn left a comment

Choose a reason for hiding this comment

ahmedabu98 commented Dec 30, 2024

ahmedabu98 commented Jan 2, 2025

Abacn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ahmedabu98 commented Nov 26, 2024 •

edited

Loading

ahmedabu98 Jan 3, 2025 •

edited

Loading

ahmedabu98 Jan 3, 2025 •

edited

Loading

ahmedabu98 Jan 3, 2025 •

edited

Loading