[BigQueryIO] fetch updated schema for newly created Storage API stream writers #33231
Conversation
ahmedabu98
changed the title
|
R: @Abacn |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
…api_schemaupdate_dynamicdest
| @@ -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( | |||
There was a problem hiding this comment.
Let's add a boolean parameter to the method, so we only return schema if requested
There was a problem hiding this comment.
FYI in our code base we always perform this call for the sole purpose of fetching the schema
| @@ -531,6 +532,30 @@ public void process( | |||
| element.getKey().getKey(), dynamicDestinations, datasetService); | |||
| tableSchema = converter.getTableSchema(); | |||
| descriptor = converter.getDescriptor(false); | |||
|
|
|||
| if (autoUpdateSchema) { | |||
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
Fixes #33238
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.