From 2b4ecc526c5c1788e8ed929b09527504279d4a82 Mon Sep 17 00:00:00 2001
From: Gary White Jr <7660110+GaryPWhite@users.noreply.github.com>
Date: Thu, 15 Oct 2026 20:01:43 -0400
Subject: [PATCH 001/108] add WRITE_TRUNCATE_DATA write_disposition to the
bigquery python sdk
---
sdks/python/apache_beam/io/gcp/bigquery.py | 5 +++++
.../clients/bigquery/bigquery_v2_messages.py | 15 +++++++++------
2 files changed, 14 insertions(+), 6 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py
index 2a52b31edf7f..150724cbca15 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -590,6 +590,7 @@ class BigQueryDisposition(object):
CREATE_NEVER = 'CREATE_NEVER'
CREATE_IF_NEEDED = 'CREATE_IF_NEEDED'
+ WRITE_TRUNCATE_DATA = 'WRITE_TRUNCATE_DATA'
WRITE_TRUNCATE = 'WRITE_TRUNCATE'
WRITE_APPEND = 'WRITE_APPEND'
WRITE_EMPTY = 'WRITE_EMPTY'
@@ -606,6 +607,7 @@ def validate_create(disposition):
@staticmethod
def validate_write(disposition):
values = (
+ BigQueryDisposition.WRITE_TRUNCATE_DATA,
BigQueryDisposition.WRITE_TRUNCATE,
BigQueryDisposition.WRITE_APPEND,
BigQueryDisposition.WRITE_EMPTY)
@@ -1517,6 +1519,7 @@ def __init__(
- BigQueryDisposition.CREATE_NEVER: fail the write if does not exist.
write_disposition: A string describing what happens if the table has
already some data. Possible values are:
+ - BigQueryDisposition.WRITE_TRUNCATE_DATA: delete table data but preserve table metadata.
- BigQueryDisposition.WRITE_TRUNCATE: delete existing rows.
- BigQueryDisposition.WRITE_APPEND: add to existing rows.
- BigQueryDisposition.WRITE_EMPTY: fail the write if table not empty.
@@ -1558,6 +1561,7 @@ def __init__(
self.create_disposition = create_disposition
self.write_disposition = write_disposition
if write_disposition in (BigQueryDisposition.WRITE_EMPTY,
+ BigQueryDisposition.WRITE_TRUNCATE_DATA,
BigQueryDisposition.WRITE_TRUNCATE):
raise ValueError(
'Write disposition %s is not supported for'
@@ -2126,6 +2130,7 @@ def __init__(
write_disposition (BigQueryDisposition): A string describing what happens
if the table has already some data. Possible values are:
+ * :attr:`BigQueryDisposition.WRITE_TRUNCATE_DATA`: delete table data but preserve table metadata.
* :attr:`BigQueryDisposition.WRITE_TRUNCATE`: delete existing rows.
* :attr:`BigQueryDisposition.WRITE_APPEND`: add to existing rows.
* :attr:`BigQueryDisposition.WRITE_EMPTY`: fail the write if table not
diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
index 3392bdaf9592..8b9fa0f883f4 100644
--- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
+++ b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
@@ -5440,12 +5440,15 @@ class JobConfigurationTableCopy(_messages.Message):
destination table already exists. The following values are supported: *
WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the
table data and uses the schema and table constraints from the source
- table. * WRITE_APPEND: If the table already exists, BigQuery appends the
- data to the table. * WRITE_EMPTY: If the table already exists and
- contains data, a 'duplicate' error is returned in the job result. The
- default value is WRITE_EMPTY. Each action is atomic and only occurs if
- BigQuery is able to complete the job successfully. Creation, truncation
- and append actions occur as one atomic update upon job completion.
+ table. * WRITE_TRUNCATE_DATA: If the table already exists, BigQuery
+ overwrites the table data, but does not use schema and table constraints
+ from the source table.* WRITE_APPEND: If the table already exists,
+ BigQuery appends the data to the table. * WRITE_EMPTY: If the table
+ already exists and contains data, a 'duplicate' error is returned in
+ the job result. The default value is WRITE_EMPTY. Each action is atomic
+ and only occurs if BigQuery is able to complete the job successfully.
+ Creation, truncation and append actions occur as one atomic update
+ upon job completion.
"""
class OperationTypeValueValuesEnum(_messages.Enum):
From 6f9745d83438242a03ee2b546554bdbf33f2e9c1 Mon Sep 17 00:00:00 2001
From: Gary White Jr <7660110+GaryPWhite@users.noreply.github.com>
Date: Thu, 15 Oct 2026 20:01:43 -0400
Subject: [PATCH 002/108] changes.md
---
CHANGES.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/CHANGES.md b/CHANGES.md
index 537c3a32447c..7b56e22c9ce6 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -66,7 +66,7 @@
* Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
* (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)).
-
+* (Python) The BigQuery IO writer now allows the `WRITE_TRUNCATE_DATA` write_disposition
## New Features / Improvements
From a5a8a77ffd1e8a7576546f652bac3ab12d5bf319 Mon Sep 17 00:00:00 2001
From: Gary White Jr <7660110+GaryPWhite@users.noreply.github.com>
Date: Thu, 15 Oct 2026 21:04:54 -0400
Subject: [PATCH 003/108] revert changes to bigquery_v2_messages
---
.../clients/bigquery/bigquery_v2_messages.py | 15 ++++++---------
1 file changed, 6 insertions(+), 9 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
index 8b9fa0f883f4..3392bdaf9592 100644
--- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
+++ b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
@@ -5440,15 +5440,12 @@ class JobConfigurationTableCopy(_messages.Message):
destination table already exists. The following values are supported: *
WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the
table data and uses the schema and table constraints from the source
- table. * WRITE_TRUNCATE_DATA: If the table already exists, BigQuery
- overwrites the table data, but does not use schema and table constraints
- from the source table.* WRITE_APPEND: If the table already exists,
- BigQuery appends the data to the table. * WRITE_EMPTY: If the table
- already exists and contains data, a 'duplicate' error is returned in
- the job result. The default value is WRITE_EMPTY. Each action is atomic
- and only occurs if BigQuery is able to complete the job successfully.
- Creation, truncation and append actions occur as one atomic update
- upon job completion.
+ table. * WRITE_APPEND: If the table already exists, BigQuery appends the
+ data to the table. * WRITE_EMPTY: If the table already exists and
+ contains data, a 'duplicate' error is returned in the job result. The
+ default value is WRITE_EMPTY. Each action is atomic and only occurs if
+ BigQuery is able to complete the job successfully. Creation, truncation
+ and append actions occur as one atomic update upon job completion.
"""
class OperationTypeValueValuesEnum(_messages.Enum):
From d233c0c51f69bb0db3942a54e92daab979836cc7 Mon Sep 17 00:00:00 2001
From: "Gary White Jr." <7660110+GaryPWhite@users.noreply.github.com>
Date: Thu, 15 Oct 2026 21:05:29 -0400
Subject: [PATCH 004/108] Update CHANGES.md
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
CHANGES.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/CHANGES.md b/CHANGES.md
index 7b56e22c9ce6..7ef1c252c060 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -66,7 +66,7 @@
* Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
* (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)).
-* (Python) The BigQuery IO writer now allows the `WRITE_TRUNCATE_DATA` write_disposition
+* (Python) The BigQuery IO writer now allows the `WRITE_TRUNCATE_DATA` write_disposition ([#40622](https://github.com/apache/beam/pull/40622)).
## New Features / Improvements
From c43e2ea3c2ee683712f3fc5ff6e907090b4401da Mon Sep 17 00:00:00 2001
From: Gary White Jr <7660110+GaryPWhite@users.noreply.github.com>
Date: Thu, 15 Oct 2026 21:08:03 -0400
Subject: [PATCH 005/108] trim doc line length
---
sdks/python/apache_beam/io/gcp/bigquery.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py
index 150724cbca15..34051e98a992 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -1519,7 +1519,7 @@ def __init__(
- BigQueryDisposition.CREATE_NEVER: fail the write if does not exist.
write_disposition: A string describing what happens if the table has
already some data. Possible values are:
- - BigQueryDisposition.WRITE_TRUNCATE_DATA: delete table data but preserve table metadata.
+ - BigQueryDisposition.WRITE_TRUNCATE_DATA: delete rows, keep schema.
- BigQueryDisposition.WRITE_TRUNCATE: delete existing rows.
- BigQueryDisposition.WRITE_APPEND: add to existing rows.
- BigQueryDisposition.WRITE_EMPTY: fail the write if table not empty.
From 9335f60e8b1bf80bf648a864e87ba4edc56c0e31 Mon Sep 17 00:00:00 2001
From: Gary White Jr <7660110+GaryPWhite@users.noreply.github.com>
Date: Thu, 15 Oct 2026 21:42:45 -0400
Subject: [PATCH 006/108] add WRITE_TRUNCATE_DATA support to bigqueryIO in java
---
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 19 +++++++++++++++--
.../gcp/bigquery/UpdateSchemaDestination.java | 3 +++
.../providers/BigQueryWriteConfiguration.java | 3 ++-
.../apache_beam/io/gcp/bigquery_file_loads.py | 21 ++++++++++++-------
.../clients/bigquery/bigquery_v2_messages.py | 15 +++++++------
5 files changed, 45 insertions(+), 16 deletions(-)
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
index 0cefea4639fd..2f1d6b702691 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -3038,6 +3038,17 @@ public enum WriteDisposition {
*/
WRITE_TRUNCATE,
+ /**
+ * Specifies that write should replace a table, but not the table schema.
+ *
+ *
The replacement may occur in multiple steps - for instance by first removing the
+ * existing table, then creating a replacement, then filling it in. This is not an atomic
+ * operation, and external programs may see the table in any of these intermediate steps.
+ *
+ *
Note: This write disposition is only supported for the FILE_LOADS write method.
+ */
+ WRITE_TRUNCATE_DATA,
+
/** Specifies that rows may be appended to an existing table. */
WRITE_APPEND,
@@ -3973,8 +3984,9 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) {
getBigLakeConfiguration() == null,
"bigLakeConfiguration is only supported when using STORAGE_WRITE_API or STORAGE_API_AT_LEAST_ONCE.");
} else {
- if (getWriteDisposition() == WriteDisposition.WRITE_TRUNCATE) {
- LOG.error("The Storage API sink does not support the WRITE_TRUNCATE write disposition.");
+ WriteDisposition disposition = getWriteDisposition()
+ if (disposition == WriteDisposition.WRITE_TRUNCATE || disposition == WriteDisposition.WRITE_TRUNCATE_DATA) {
+ LOG.error("The Storage API sink does not support the "+disposition.toString()+" write disposition.");
}
Map bigLakeConfiguration = getBigLakeConfiguration();
if (bigLakeConfiguration != null) {
@@ -4229,6 +4241,9 @@ private WriteResult continueExpandTyped(
checkArgument(
getWriteDisposition() != WriteDisposition.WRITE_TRUNCATE,
"WriteDisposition.WRITE_TRUNCATE is not supported for an unbounded PCollection.");
+ checkArgument(
+ getWriteDisposition() != WriteDisposition.WRITE_TRUNCATE_DATA,
+ "WriteDisposition.WRITE_TRUNCATE_DATA is not supported for an unbounded PCollection.");
InsertRetryPolicy retryPolicy =
MoreObjects.firstNonNull(getFailedInsertRetryPolicy(), InsertRetryPolicy.alwaysRetry());
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java
index 536d318f0213..7f801ef38824 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java
@@ -257,6 +257,9 @@ public void finishBundle(FinishBundleContext context) throws Exception {
if (!loadConfig
.getWriteDisposition()
.equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE.toString())
+ && !loadConfig
+ .getWriteDisposition()
+ .equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE_DATA.toString()
&& !loadConfig
.getWriteDisposition()
.equals(BigQueryIO.Write.WriteDisposition.WRITE_APPEND.toString())) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java
index ec59db14501b..ec498c66aacc 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java
@@ -138,7 +138,8 @@ public static Builder builder() {
@SchemaFieldDescription(
"Specifies the action that occurs if the destination table already exists. "
+ "The following values are supported: "
- + "WRITE_TRUNCATE (overwrites the table data), "
+ + "WRITE_TRUNCATE (overwrites the table data and schema), "
+ + "WRITE_TRUNCATE_DATA (overwrites the table data), "
+ "WRITE_APPEND (append the data to the table), "
+ "WRITE_EMPTY (job must fail if the table is not empty).")
@Pure
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
index dabe80c5c27e..db51ecf22403 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
@@ -416,7 +416,11 @@ def process(self, element, schema_mod_job_name_prefix):
# When writing to normal tables WRITE_TRUNCATE will overwrite the schema but
# when writing to a partition, care needs to be taken to update the schema
# even on WRITE_TRUNCATE.
- if (self._write_disposition not in ('WRITE_TRUNCATE', 'WRITE_APPEND') or
+ if (self._write_disposition not in (
+ 'WRITE_TRUNCATE',
+ 'WRITE_TRUNCATE_DATA',
+ 'WRITE_APPEND'
+ ) or
not additional_parameters or
not additional_parameters.get("schemaUpdateOptions")):
# No need to modify schema of destination table
@@ -613,11 +617,12 @@ def _determine_write_disposition(self, copy_to_reference) -> tuple[bool, str]:
Determines the write disposition for a BigQuery copy job,
based on destination.
- When the write_disposition for a job is WRITE_TRUNCATE, multiple copy jobs
- to the same destination can interfere with each other, truncate data, and
- write to the BigQuery table repeatedly. To prevent this, the first copy job
- runs with the user's specified write_disposition, but subsequent jobs must
- always use WRITE_APPEND. This ensures that subsequent copy jobs do not
+ When the write_disposition for a job is WRITE_TRUNCATE or
+ WRITE_TRUNCATE_DATA, multiple copy jobs to the same destination can
+ interfere with each other, truncate data, and write to the BigQuery
+ table repeatedly. To prevent this, the first copy job runs with the
+ user's specified write_disposition, but subsequent jobs must always
+ use WRITE_APPEND. This ensures that subsequent copy jobs do not
clear out data appended by previous jobs.
Args:
@@ -1231,7 +1236,9 @@ def _load_data(
load_job_project_id=self.load_job_project_id),
schema_mod_job_name_pcv))
- if self.write_disposition in ('WRITE_EMPTY', 'WRITE_TRUNCATE'):
+ if self.write_disposition in (
+ 'WRITE_EMPTY', 'WRITE_TRUNCATE', 'WRITE_TRUNCATE_DATA'
+ ):
# All loads going to the same table must be processed together so that
# the truncation happens only once. See
# https://github.com/apache/beam/issues/24535.
diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
index 3392bdaf9592..8b9fa0f883f4 100644
--- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
+++ b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
@@ -5440,12 +5440,15 @@ class JobConfigurationTableCopy(_messages.Message):
destination table already exists. The following values are supported: *
WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the
table data and uses the schema and table constraints from the source
- table. * WRITE_APPEND: If the table already exists, BigQuery appends the
- data to the table. * WRITE_EMPTY: If the table already exists and
- contains data, a 'duplicate' error is returned in the job result. The
- default value is WRITE_EMPTY. Each action is atomic and only occurs if
- BigQuery is able to complete the job successfully. Creation, truncation
- and append actions occur as one atomic update upon job completion.
+ table. * WRITE_TRUNCATE_DATA: If the table already exists, BigQuery
+ overwrites the table data, but does not use schema and table constraints
+ from the source table.* WRITE_APPEND: If the table already exists,
+ BigQuery appends the data to the table. * WRITE_EMPTY: If the table
+ already exists and contains data, a 'duplicate' error is returned in
+ the job result. The default value is WRITE_EMPTY. Each action is atomic
+ and only occurs if BigQuery is able to complete the job successfully.
+ Creation, truncation and append actions occur as one atomic update
+ upon job completion.
"""
class OperationTypeValueValuesEnum(_messages.Enum):
From 86504f71ce90914c30846e90cd4dccaea41b5a67 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Fri, 16 Oct 2026 07:34:15 -0400
Subject: [PATCH 007/108] Bump google.golang.org/api from 0.301.0 to 0.302.0 in
/sdks (#40627)
Bumps [google.golang.org/api](https://github.com/googleapis/google-api-go-client) from 0.301.0 to 0.302.0.
- [Release notes](https://github.com/googleapis/google-api-go-client/releases)
- [Changelog](https://github.com/googleapis/google-api-go-client/blob/main/CHANGES.md)
- [Commits](https://github.com/googleapis/google-api-go-client/compare/v0.301.0...v0.302.0)
---
updated-dependencies:
- dependency-name: google.golang.org/api
dependency-version: 0.302.0
dependency-type: direct:production
update-type: version-update:semver-minor
...
Signed-off-by: dependabot[bot]
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
sdks/go.mod | 8 ++++----
sdks/go.sum | 16 ++++++++--------
2 files changed, 12 insertions(+), 12 deletions(-)
diff --git a/sdks/go.mod b/sdks/go.mod
index e59d8b8c7a40..d60a9db1651f 100644
--- a/sdks/go.mod
+++ b/sdks/go.mod
@@ -60,7 +60,7 @@ require (
golang.org/x/sync v0.24.0
golang.org/x/sys v0.49.0
golang.org/x/text v0.43.0
- google.golang.org/api v0.301.0
+ google.golang.org/api v0.302.0
google.golang.org/genproto v0.0.0-20260715232425-e75dac1f907d // indirect
google.golang.org/grpc v1.84.0
google.golang.org/protobuf v1.36.12
@@ -171,7 +171,7 @@ require (
github.com/google/pprof v0.0.0-20260604005048-7023385849c0 // indirect
github.com/google/renameio/v2 v2.0.0 // indirect
github.com/google/s2a-go v0.1.11 // indirect
- github.com/googleapis/enterprise-certificate-proxy v0.3.22 // indirect
+ github.com/googleapis/enterprise-certificate-proxy v0.3.23 // indirect
github.com/googleapis/gax-go/v2 v2.26.2 // indirect
github.com/gorilla/handlers v1.5.2 // indirect
github.com/gorilla/mux v1.8.1 // indirect
@@ -198,9 +198,9 @@ require (
go.opencensus.io v0.24.0 // indirect
go.yaml.in/yaml/v3 v3.0.5 // indirect
golang.org/x/crypto v0.58.0 // indirect
- golang.org/x/mod v0.41.0 // indirect
+ golang.org/x/mod v0.42.0 // indirect
golang.org/x/tools v0.51.0 // indirect
golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20260715232425-e75dac1f907d // indirect
- google.golang.org/genproto/googleapis/rpc v0.0.0-20260928230214-8a89bd6388cc // indirect
+ google.golang.org/genproto/googleapis/rpc v0.0.0-20261005182115-fad411399dd8 // indirect
)
diff --git a/sdks/go.sum b/sdks/go.sum
index 0bd5914755f8..1fc97c94e394 100644
--- a/sdks/go.sum
+++ b/sdks/go.sum
@@ -533,8 +533,8 @@ github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/wire v0.5.0/go.mod h1:ngWDr9Qvq3yZA10YrxfyGELY/AFWGVpy9c1LTRi1EoU=
-github.com/googleapis/enterprise-certificate-proxy v0.3.22 h1:NU4XpII6jD+Dxcot94fqjE+AfJoE/lQP9q3faYGzC/c=
-github.com/googleapis/enterprise-certificate-proxy v0.3.22/go.mod h1:L3D/IQExI6LqEjBdXcZQ1WluSgigQmSwBboFstVPM4w=
+github.com/googleapis/enterprise-certificate-proxy v0.3.23 h1:J+VGF1wjAM2dU2kKUXZxJdDppRLReWP7KWcclfngQhc=
+github.com/googleapis/enterprise-certificate-proxy v0.3.23/go.mod h1:Y8NO/15GNLkrXI3OR5xOFmeU3XLX1sLSdtlXgJ0zmIw=
github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
github.com/googleapis/gax-go/v2 v2.1.0/go.mod h1:Q3nei7sK6ybPYH7twZdmQpAd1MKb7pfu6SK+H1/DsU0=
@@ -976,8 +976,8 @@ golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/mod v0.5.0/go.mod h1:5OXOZSfqPIIbmVBIIKWRFfZjPR0E5r58TLhUjH0a2Ro=
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
-golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c=
-golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o=
+golang.org/x/mod v0.42.0 h1:ICeL3t2L30ljSPHHxDnR4fqVcBt/zR2zJAb9juavmQg=
+golang.org/x/mod v0.42.0/go.mod h1:u5fzIrRInkL8MaO1cbRkUHT64v303YIT5KEqpNr7jRs=
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20190108225652-1e06a53dbb7e/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -1315,8 +1315,8 @@ google.golang.org/api v0.69.0/go.mod h1:boanBiw+h5c3s+tBPgEzLDRHfFLWV0qXxRHz3ws7
google.golang.org/api v0.70.0/go.mod h1:Bs4ZM2HGifEvXwd50TtW70ovgJffJYw2oRCOFU/SkfA=
google.golang.org/api v0.71.0/go.mod h1:4PyU6e6JogV1f9eA4voyrTY2batOLdgZ5qZ5HOCc4j8=
google.golang.org/api v0.74.0/go.mod h1:ZpfMZOVRMywNyvJFeqL9HRWBgAuRfSjJFpe9QtRRyDs=
-google.golang.org/api v0.301.0 h1:rTvI8szs17R7b09ggb34Bv4XrmFaFl9uCBncgxskqtw=
-google.golang.org/api v0.301.0/go.mod h1:gGUTaOjgsaG/KNyyW1tFO+DSFm3EJw4BwyehlcmDCXs=
+google.golang.org/api v0.302.0 h1:/eENsEv3+GE0QiEmXjdaLI84Mt76LNjGz4vp4E7sRVw=
+google.golang.org/api v0.302.0/go.mod h1:dcjRWAy1QdjGb3Z67Qu0gdBWH2gL8Qoy/RI2ZVuvdeE=
google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
google.golang.org/appengine v1.5.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
@@ -1414,8 +1414,8 @@ google.golang.org/genproto v0.0.0-20260715232425-e75dac1f907d h1:C9v1o0/4quuhOAf
google.golang.org/genproto v0.0.0-20260715232425-e75dac1f907d/go.mod h1:Wz2wFJntZFmLGo7pLDXZ3wYk5hyc0Mb+SkHhDDXT+lU=
google.golang.org/genproto/googleapis/api v0.0.0-20260715232425-e75dac1f907d h1:QwnJwPte4XXAkhPu26LTDIahnsMSUV0kK8HkxbC+Pc4=
google.golang.org/genproto/googleapis/api v0.0.0-20260715232425-e75dac1f907d/go.mod h1:WRrQ7/7N19PypuT0fxLOL5Lq0waoiRri4FbtHDEKrGE=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260928230214-8a89bd6388cc h1:4bNTbnb44EqGVy9HaQxSY2jnafSUdgV3qHtedvNpDKg=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260928230214-8a89bd6388cc/go.mod h1:OaIUM3+LpYcK2GXM4FTmhWoIq371Owdr+Cc7/BsYHHc=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20261005182115-fad411399dd8 h1:0SXBqli5dnuCKMjI/X3FM226zEaEa0GUcSCJtN1vL2o=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20261005182115-fad411399dd8/go.mod h1:OaIUM3+LpYcK2GXM4FTmhWoIq371Owdr+Cc7/BsYHHc=
google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=
From 0072521920c2d4219f3191e2ffd748ee4b84babb Mon Sep 17 00:00:00 2001
From: Shunping Huang
Date: Fri, 16 Oct 2026 08:41:32 -0400
Subject: [PATCH 008/108] [Java] Fix GRPC channel leak (#39663)
---
.../fn/harness/ExternalWorkerService.java | 33 +++++++++++--------
.../org/apache/beam/fn/harness/FnHarness.java | 1 +
.../harness/control/ProcessBundleHandler.java | 3 ++
.../fn/harness/data/BeamFnDataGrpcClient.java | 11 +++++--
.../state/BeamFnStateGrpcClientCache.java | 20 +++++++++--
5 files changed, 50 insertions(+), 18 deletions(-)
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java
index 466d8ef3193b..b1bc462702df 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ExternalWorkerService.java
@@ -39,6 +39,7 @@
import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.sdk.util.construction.Environments;
import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel;
import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
@@ -78,20 +79,24 @@ public void startWorker(
.withInterceptors(
ImmutableList.of(AddHarnessIdInterceptor.create(request.getWorkerId())));
- ProvisionServiceGrpc.ProvisionServiceBlockingStub provisionStub =
- ProvisionServiceGrpc.newBlockingStub(
- channelFactory.forDescriptor(request.getProvisionEndpoint()));
- ProvisionApi.ProvisionInfo provisionInfo =
- provisionStub
- .getProvisionInfo(ProvisionApi.GetProvisionInfoRequest.newBuilder().build())
- .getInfo();
-
- runnerCapabilites = Sets.newHashSet(provisionInfo.getRunnerCapabilitiesList());
- if (provisionInfo.hasControlEndpoint()) {
- controlEndpoint = provisionInfo.getControlEndpoint();
- }
- if (provisionInfo.hasLoggingEndpoint()) {
- loggingEndpoint = provisionInfo.getLoggingEndpoint();
+ ManagedChannel channel = channelFactory.forDescriptor(request.getProvisionEndpoint());
+ try {
+ ProvisionServiceGrpc.ProvisionServiceBlockingStub provisionStub =
+ ProvisionServiceGrpc.newBlockingStub(channel);
+ ProvisionApi.ProvisionInfo provisionInfo =
+ provisionStub
+ .getProvisionInfo(ProvisionApi.GetProvisionInfoRequest.newBuilder().build())
+ .getInfo();
+
+ runnerCapabilites = Sets.newHashSet(provisionInfo.getRunnerCapabilitiesList());
+ if (provisionInfo.hasControlEndpoint()) {
+ controlEndpoint = provisionInfo.getControlEndpoint();
+ }
+ if (provisionInfo.hasLoggingEndpoint()) {
+ loggingEndpoint = provisionInfo.getLoggingEndpoint();
+ }
+ } finally {
+ channel.shutdown();
}
}
// Lambda closured variables must be final.
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
index 63f6c41b5a71..f547a1eeadff 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
@@ -437,6 +437,7 @@ private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) {
beamFnStatusClient.close();
}
processBundleHandler.shutdown();
+ channel.shutdown();
} catch (Exception e) {
LOG.error("Shutting down harness due to exception", e);
e.printStackTrace();
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
index f12340552fd0..9fd1f0f0d68b 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
@@ -768,6 +768,9 @@ public BeamFnApi.InstructionResponse.Builder trySplit(InstructionRequest request
public void shutdown() throws Exception {
bundleProcessorCache.shutdown();
beamFnDataClient.close();
+ if (beamFnStateGrpcClientCache != null) {
+ beamFnStateGrpcClientCache.close();
+ }
}
@VisibleForTesting
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
index 79e34fe5765f..0640cc0d289d 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
@@ -76,6 +76,7 @@ public int hashCode() {
}
private final ConcurrentMap multiplexerCache;
+ private final ConcurrentMap channelCache;
private final Function channelFactory;
private final OutboundObserverFactory outboundObserverFactory;
@@ -85,6 +86,7 @@ public BeamFnDataGrpcClient(
this.channelFactory = channelFactory;
this.outboundObserverFactory = outboundObserverFactory;
this.multiplexerCache = new ConcurrentHashMap<>();
+ this.channelCache = new ConcurrentHashMap<>();
}
@Override
@@ -128,6 +130,10 @@ public void close() {
}
}
multiplexerCache.clear();
+ for (ManagedChannel channel : channelCache.values()) {
+ channel.shutdown();
+ }
+ channelCache.clear();
}
@Override
@@ -142,10 +148,11 @@ private BeamFnDataGrpcMultiplexer getMultiplexer(
return multiplexerCache.computeIfAbsent(
key,
k -> {
+ ManagedChannel channel =
+ channelCache.computeIfAbsent(k, unused -> channelFactory.apply(apiServiceDescriptor));
OutboundObserverFactory.BasicFactory baseOutboundObserverFactory =
inboundObserver -> {
- BeamFnDataGrpc.BeamFnDataStub stub =
- BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
+ BeamFnDataGrpc.BeamFnDataStub stub = BeamFnDataGrpc.newStub(channel);
if (dataStreamId != null && !dataStreamId.isEmpty()) {
Metadata headers = new Metadata();
headers.put(
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
index c9adb9d7d142..1a0c4f3fb912 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
@@ -18,7 +18,9 @@
package org.apache.beam.fn.harness.state;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
@@ -44,7 +46,7 @@
public class BeamFnStateGrpcClientCache {
private static final Logger LOG = LoggerFactory.getLogger(BeamFnStateGrpcClientCache.class);
- private final Map cache;
+ private final Map cache;
private final ManagedChannelFactory channelFactory;
private final OutboundObserverFactory outboundObserverFactory;
private final IdGenerator idGenerator;
@@ -71,7 +73,7 @@ public synchronized BeamFnStateClient forApiServiceDescriptor(
// We specifically are synchronized so that we only create one GrpcStateClient at a time
// preventing a race where multiple GrpcStateClient objects might be constructed at the same
// for the same ApiServiceDescriptor.
- BeamFnStateClient rval;
+ GrpcStateClient rval;
synchronized (cache) {
rval = cache.get(apiServiceDescriptor);
}
@@ -87,6 +89,18 @@ public synchronized BeamFnStateClient forApiServiceDescriptor(
return rval;
}
+ /** Shuts down the underlying channels of all cached state clients. */
+ public void close() {
+ List clients;
+ synchronized (cache) {
+ clients = new ArrayList<>(cache.values());
+ cache.clear();
+ }
+ for (GrpcStateClient client : clients) {
+ client.channel.shutdown();
+ }
+ }
+
/** A {@link BeamFnStateClient} for a given {@link ApiServiceDescriptor}. */
private class GrpcStateClient implements BeamFnStateClient {
private final Object lock = new Object();
@@ -111,6 +125,7 @@ private GrpcStateClient(ApiServiceDescriptor apiServiceDescriptor) {
synchronized (lock) {
if (errorDuringConstruction) {
outboundObserver.onCompleted();
+ channel.shutdown();
}
}
}
@@ -160,6 +175,7 @@ private void closeAndCleanUp(RuntimeException cause) {
errorDuringConstruction = true;
} else {
outboundObserver.onCompleted();
+ channel.shutdown();
}
}
}
From 3336b153784ce57ef226acfa6efafd8feccf2bd8 Mon Sep 17 00:00:00 2001
From: pablo rodriguez defino
Date: Fri, 16 Oct 2026 07:11:09 -0700
Subject: [PATCH 009/108] [Prism] Guard ClearOrderedListState against
uninitialized state maps (#40626)
* [Prism] Guard ClearOrderedListState against uninitialized state maps
Problem:
In TentativeData.ClearOrderedListState (sdks/go/pkg/beam/runners/prism/internal/engine/data.go),
Prism indexed directly into d.state[stateID][w] and wrote the cleared slice back to
kMap[string(uKey)] without checking whether d.state, d.state[stateID], or winMap[w]
had been initialized:
winMap := d.state[stateID]
w := d.toWindow(wKey)
kMap := winMap[w]
data := kMap[string(uKey)]
...
kMap[string(uKey)] = StateData{Bag: cleared} // panic: assignment to entry in nil map
SDK harnesses (such as the Python SDK's SynchronousOrderedListRuntimeState.commit and
the Rust SDK's OrderedListCell) buffer ordered-list mutations in memory during a bundle
and flush pending StateRequest_Clear range requests before StateRequest_Append requests
at bundle commit. When a bundle clears an ordered-list state cell (or range) in a
window for which no prior bundle has persisted state on the runner, TentativeData.state
has no map for (stateID, window), kMap is nil, and Prism panics on line 336.
Fix:
Use the existing d.clearState(stateID, wKey) helper in ClearOrderedListState and return
early when kMap is nil, matching ClearBagState, ClearMultimapState, and
ClearMultimapKeysState.
Reproduction (Python SDK on PrismRunner):
```python
import apache_beam as beam
from apache_beam.coders import VarIntCoder
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.testing.util import assert_that, equal_to
from apache_beam.transforms.userstate import OrderedListStateSpec
from apache_beam.utils.timestamp import Timestamp
class ClearOrderedListDoFn(beam.DoFn):
EVENTS = OrderedListStateSpec("events", VarIntCoder())
def process(self, element, events=beam.DoFn.StateParam(EVENTS)):
key, (ts, op) = element
if op == "CLEAR":
events.clear()
yield f"cleared:{key}"
else:
events.add((Timestamp(micros=ts), op))
entries = [f"{int(t.micros)}={v}" for t, v in events.read()]
yield f"entry:{key}:{','.join(entries)}"
if __name__ == "__main__":
options = PipelineOptions(["--runner=PrismRunner", "--environment_type=LOOPBACK"])
with beam.Pipeline(options=options) as p:
out = (
p
| beam.Create([("k", (30, 300)), ("k", (10, 100)), ("k", (0, "CLEAR"))])
| beam.ParDo(ClearOrderedListDoFn())
)
assert_that(
out,
equal_to(["entry:k:30=300", "entry:k:10=100,30=300", "cleared:k"]),
)
```
* Refactor comments in data_test.go
Updated comments for clarity in data_test.go.
---
.../runners/prism/internal/engine/data.go | 7 ++--
.../prism/internal/engine/data_test.go | 37 +++++++++++++++++++
2 files changed, 41 insertions(+), 3 deletions(-)
diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/data.go b/sdks/go/pkg/beam/runners/prism/internal/engine/data.go
index 371e3f00b895..23553413f44f 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/engine/data.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/engine/data.go
@@ -322,9 +322,10 @@ func findRange(bag [][]byte, start, end int64) (int, int) {
}
func (d *TentativeData) ClearOrderedListState(stateID LinkID, wKey, uKey []byte, start, end int64) {
- winMap := d.state[stateID]
- w := d.toWindow(wKey)
- kMap := winMap[w]
+ kMap := d.clearState(stateID, wKey)
+ if kMap == nil {
+ return
+ }
data := kMap[string(uKey)]
lo, hi := findRange(data.Bag, start, end)
diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/data_test.go b/sdks/go/pkg/beam/runners/prism/internal/engine/data_test.go
index 1d0497104182..8d800e3cd5f1 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/engine/data_test.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/engine/data_test.go
@@ -219,4 +219,41 @@ func TestOrderedListState(t *testing.T) {
t.Errorf("OrderedList int32 \n%v", d)
}
})
+ t.Run("clear_uninitialized", func(t *testing.T) {
+ d := TentativeData{
+ stateTypeLen: map[LinkID]func([]byte) int{
+ linkID: func(_ []byte) int {
+ return 1
+ },
+ },
+ }
+
+ // d.state is nil.
+ d.ClearOrderedListState(linkID, wKey, uKey, math.MinInt64, math.MaxInt64)
+ if d.state != nil {
+ t.Fatalf("ClearOrderedListState on empty TentativeData allocated state map: %v", d.state)
+ }
+
+ // d.state is non-nil, but has no entry for linkID.
+ otherLink := LinkID{Transform: "dofn", Local: "otherState"}
+ d.AppendBagState(otherLink, wKey, uKey, []byte{1})
+ d.ClearOrderedListState(linkID, wKey, uKey, 0, 100)
+ if _, ok := d.state[linkID]; ok {
+ t.Fatalf("ClearOrderedListState created entry for uninitialized stateID %v", linkID)
+ }
+
+ // d.state[linkID] exists for an interval window, but not for the global window.
+ // An SDK that flushes range clears before appends on the first bundle in a window
+ // must not panic and must preserve subsequent appends.
+ iwKey := make([]byte, 16)
+ binary.BigEndian.PutUint64(iwKey[0:8], uint64(1000)^0x8000000000000000)
+ d.AppendOrderedListState(linkID, iwKey, uKey, cc(time1, 1))
+ d.ClearOrderedListState(linkID, wKey, uKey, 0, 30)
+ d.AppendOrderedListState(linkID, wKey, uKey, cc(time2, 1))
+ got := d.GetOrderedListState(linkID, wKey, uKey, 0, 60)
+ want := [][]byte{cc(time2, 1)}
+ if diff := cmp.Diff(want, got); diff != "" {
+ t.Errorf("OrderedList after clear_uninitialized + append:\n%v", diff)
+ }
+ })
}
From f3a5f4f15425fb1979e8843ab089477f1274bf98 Mon Sep 17 00:00:00 2001
From: Goutam Adwant <8672451+goutamadwant@users.noreply.github.com>
Date: Fri, 16 Oct 2026 07:33:43 -0700
Subject: [PATCH 010/108] Add stale GCS test bucket cleaner (#40386)
---
.test-infra/tools/README.md | 18 ++
.test-infra/tools/build.gradle | 10 +
.../tools/stale_gcs_buckets_cleaner.py | 112 ++++++++++
.../tools/test_stale_gcs_buckets_cleaner.py | 202 ++++++++++++++++++
4 files changed, 342 insertions(+)
create mode 100644 .test-infra/tools/stale_gcs_buckets_cleaner.py
create mode 100644 .test-infra/tools/test_stale_gcs_buckets_cleaner.py
diff --git a/.test-infra/tools/README.md b/.test-infra/tools/README.md
index 758c3885a814..f13b672ddde8 100644
--- a/.test-infra/tools/README.md
+++ b/.test-infra/tools/README.md
@@ -22,6 +22,24 @@
## Tools
- [Python installer](#python-installer)
+- [Stale GCS test bucket cleaner](#stale-gcs-test-bucket-cleaner)
+
+### Stale GCS test bucket cleaner
+
+`stale_gcs_buckets_cleaner.py` deletes buckets in the `apache-beam-testing`
+project whose names match `apache-beam-temp-bucket-` and are older than
+24 hours. These buckets are created by `GcsUtil` integration tests and can be
+left behind when a test is interrupted before teardown. The cleaner removes
+objects from eligible buckets before deleting them.
+
+The cleaner performs a dry run by default:
+
+```bash
+python3 stale_gcs_buckets_cleaner.py
+```
+
+The scheduled `Clean Up GCP Resources` workflow runs the cleaner through the
+`removeStaleGcsBuckets` Gradle task, which passes the explicit `--delete` flag.
### Python installer
diff --git a/.test-infra/tools/build.gradle b/.test-infra/tools/build.gradle
index 92b53f2fd560..8c61e357c042 100644
--- a/.test-infra/tools/build.gradle
+++ b/.test-infra/tools/build.gradle
@@ -40,6 +40,15 @@ task removeStaleSpannerResources(type: Exec) {
commandLine './stale_spanner_cleaner.sh'
}
+task testStaleGcsBucketsCleaner(type: Exec) {
+ commandLine 'bash', '-c', 'python3 -m pip install -r requirements.txt && python3 -m unittest test_stale_gcs_buckets_cleaner.py'
+}
+
+task removeStaleGcsBuckets(type: Exec) {
+ dependsOn tasks.testStaleGcsBucketsCleaner
+ commandLine 'python3', 'stale_gcs_buckets_cleaner.py', '--delete'
+}
+
task runStaleCleaner(type: Exec) {
commandLine 'bash', '-c', 'pip install -r requirements.txt && python3 stale_cleaner.py'
}
@@ -50,5 +59,6 @@ task cleanupOtherStaleResources {
finalizedBy tasks.removeStaleCbtInstances
finalizedBy tasks.removeStaleK8sWorkload
finalizedBy tasks.removeStaleSpannerResources
+ finalizedBy tasks.removeStaleGcsBuckets
finalizedBy tasks.runStaleCleaner
}
diff --git a/.test-infra/tools/stale_gcs_buckets_cleaner.py b/.test-infra/tools/stale_gcs_buckets_cleaner.py
new file mode 100644
index 000000000000..4d475018862e
--- /dev/null
+++ b/.test-infra/tools/stale_gcs_buckets_cleaner.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+"""Deletes stale GCS buckets left behind by GcsUtil integration tests."""
+
+import argparse
+import datetime
+import re
+
+from google.cloud import storage
+
+
+DEFAULT_PROJECT_ID = "apache-beam-testing"
+GCS_TEMP_BUCKET_PREFIX = "apache-beam-temp-bucket-"
+GCS_TEMP_BUCKET_PATTERN = re.compile(
+ rf"^{GCS_TEMP_BUCKET_PREFIX}"
+ r"[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-"
+ r"[0-9a-f]{12}$")
+DEFAULT_MAX_AGE = datetime.timedelta(hours=24)
+UTC = datetime.timezone.utc
+
+
+def _as_utc(value):
+ if value.tzinfo is None:
+ return value.replace(tzinfo=UTC)
+ return value.astimezone(UTC)
+
+
+def clean_stale_gcs_buckets(
+ client,
+ max_age=DEFAULT_MAX_AGE,
+ now=None,
+ dry_run=True,
+):
+ """Deletes test buckets whose creation time is older than ``max_age``.
+
+ Returns the names of buckets selected for deletion. If a deletion fails,
+ the cleaner attempts the remaining buckets before reporting the failures.
+ """
+ current_time = _as_utc(now or datetime.datetime.now(UTC))
+ selected = []
+ failures = []
+
+ buckets = client.list_buckets(
+ project=DEFAULT_PROJECT_ID, prefix=GCS_TEMP_BUCKET_PREFIX)
+ for bucket in buckets:
+ # Keep this check even though list_buckets filters server-side. It is
+ # the final guard before a destructive operation.
+ if not GCS_TEMP_BUCKET_PATTERN.fullmatch(bucket.name):
+ continue
+ if bucket.time_created is None:
+ print(f"Skipping {bucket.name}: creation time is unavailable")
+ continue
+ if current_time - _as_utc(bucket.time_created) <= max_age:
+ continue
+
+ selected.append(bucket.name)
+ if dry_run:
+ print(f"Dry run: would delete gs://{bucket.name}")
+ continue
+
+ print(f"Deleting stale test bucket gs://{bucket.name}")
+ try:
+ # Bucket.delete(force=True) refuses to empty more than 256 objects.
+ # Delete each listed generation, then delete the empty bucket.
+ for blob in bucket.list_blobs(versions=True):
+ blob.delete(if_generation_match=blob.generation)
+ bucket.delete()
+ except Exception as error:
+ failures.append(bucket.name)
+ print(f"Failed to delete gs://{bucket.name}: {error}")
+
+ if failures:
+ raise RuntimeError(f"Bucket failure count: {len(failures)}")
+
+ return selected
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="Delete stale buckets created by GcsUtil integration tests.")
+ parser.add_argument(
+ "--delete",
+ action="store_true",
+ help="Delete eligible buckets instead of performing a dry run.",
+ )
+ args = parser.parse_args()
+
+ client = storage.Client(project=DEFAULT_PROJECT_ID)
+ selected = clean_stale_gcs_buckets(client, dry_run=not args.delete)
+ action = "Deleted" if args.delete else "Selected"
+ print(f"{action} {len(selected)} stale GCS test bucket(s)")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/.test-infra/tools/test_stale_gcs_buckets_cleaner.py b/.test-infra/tools/test_stale_gcs_buckets_cleaner.py
new file mode 100644
index 000000000000..21cbaa1dd76e
--- /dev/null
+++ b/.test-infra/tools/test_stale_gcs_buckets_cleaner.py
@@ -0,0 +1,202 @@
+#!/usr/bin/env python
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import datetime
+import unittest
+from unittest import mock
+
+import stale_gcs_buckets_cleaner as cleaner
+
+
+UTC = datetime.timezone.utc
+NOW = datetime.datetime(2026, 9, 30, 12, tzinfo=UTC)
+STALE_BUCKET = "apache-beam-temp-bucket-123e4567-e89b-42d3-a456-426614174000"
+FRESH_BUCKET = "apache-beam-temp-bucket-123e4567-e89b-42d3-a456-426614174001"
+
+
+class FakeBlob:
+ def __init__(self, bucket, name, generation):
+ self.bucket = bucket
+ self.name = name
+ self.generation = generation
+
+ def delete(self, if_generation_match=None):
+ if if_generation_match != self.generation:
+ raise ValueError("blob generation precondition is required")
+ self.bucket.blobs.remove(self)
+
+
+class FakeBucket:
+ def __init__(self, name, created, blob_count=0):
+ self.name = name
+ self.time_created = created
+ self.blobs = [FakeBlob(self, f"blob-{i}", i + 1)
+ for i in range(blob_count)]
+ self.deleted = False
+ self.delete_error = None
+
+ def list_blobs(self, versions=False):
+ if not versions:
+ raise ValueError("all object generations must be listed")
+ return iter(list(self.blobs))
+
+ def delete(self):
+ if self.delete_error:
+ raise self.delete_error
+ if self.blobs:
+ raise ValueError("bucket is not empty")
+ self.deleted = True
+
+
+class FakeClient:
+ def __init__(self):
+ self.buckets = []
+
+ def list_buckets(self, project, prefix):
+ if project != "apache-beam-testing" or prefix != "apache-beam-temp-bucket-":
+ raise ValueError("unexpected bucket selection")
+ return iter(self.buckets)
+
+
+def bucket(name, created, blob_count=0):
+ return FakeBucket(name, created, blob_count)
+
+
+class StaleGcsBucketsCleanerTest(unittest.TestCase):
+ def setUp(self):
+ self.client = FakeClient()
+
+ def test_deletes_only_stale_test_buckets(self):
+ stale = bucket(
+ STALE_BUCKET, NOW - datetime.timedelta(hours=25))
+ fresh = bucket(
+ FRESH_BUCKET, NOW - datetime.timedelta(hours=23))
+ unrelated = bucket("other-bucket", NOW - datetime.timedelta(days=7))
+ missing_created = bucket(
+ "apache-beam-temp-bucket-123e4567-e89b-42d3-a456-426614174002",
+ None)
+ self.client.buckets = [
+ stale, fresh, unrelated, missing_created]
+
+ deleted = cleaner.clean_stale_gcs_buckets(
+ self.client, now=NOW, dry_run=False)
+
+ self.assertEqual(deleted, [stale.name])
+ self.assertTrue(stale.deleted)
+ self.assertFalse(fresh.deleted)
+ self.assertFalse(unrelated.deleted)
+ self.assertFalse(missing_created.deleted)
+
+ def test_preserves_bucket_at_age_threshold(self):
+ value = bucket(
+ STALE_BUCKET,
+ NOW - datetime.timedelta(hours=24))
+ self.client.buckets = [value]
+
+ deleted = cleaner.clean_stale_gcs_buckets(
+ self.client, now=NOW, dry_run=False)
+
+ self.assertEqual(deleted, [])
+ self.assertFalse(value.deleted)
+
+ def test_dry_run_does_not_delete_bucket(self):
+ value = bucket(
+ STALE_BUCKET, NOW - datetime.timedelta(days=2))
+ self.client.buckets = [value]
+
+ selected = cleaner.clean_stale_gcs_buckets(self.client, now=NOW)
+
+ self.assertEqual(selected, [value.name])
+ self.assertFalse(value.deleted)
+
+ def test_naive_creation_time_is_treated_as_utc(self):
+ value = bucket(
+ STALE_BUCKET,
+ datetime.datetime(2026, 9, 29, 11))
+ self.client.buckets = [value]
+
+ deleted = cleaner.clean_stale_gcs_buckets(
+ self.client, now=NOW, dry_run=False)
+
+ self.assertEqual(deleted, [value.name])
+ self.assertTrue(value.deleted)
+
+ def test_continues_after_delete_failure_and_reports_all_failures(self):
+ failed = bucket(
+ STALE_BUCKET, NOW - datetime.timedelta(days=2))
+ deleted = bucket(
+ FRESH_BUCKET, NOW - datetime.timedelta(days=2))
+ failed.delete_error = RuntimeError("permission denied")
+ self.client.buckets = [failed, deleted]
+
+ with self.assertRaisesRegex(
+ RuntimeError, "Bucket failure count: 1"):
+ cleaner.clean_stale_gcs_buckets(
+ self.client, now=NOW, dry_run=False)
+
+ self.assertFalse(failed.deleted)
+ self.assertTrue(deleted.deleted)
+
+ def test_deletes_bucket_with_more_than_256_objects(self):
+ value = bucket(STALE_BUCKET, NOW - datetime.timedelta(days=2), 300)
+ self.client.buckets = [value]
+
+ selected = cleaner.clean_stale_gcs_buckets(
+ self.client, now=NOW, dry_run=False)
+
+ self.assertEqual(selected, [value.name])
+ self.assertTrue(value.deleted)
+ self.assertEqual(value.blobs, [])
+
+ def test_preserves_prefixed_bucket_without_uuid_name(self):
+ value = bucket(
+ "apache-beam-temp-bucket-manual", NOW - datetime.timedelta(days=7))
+ self.client.buckets = [value]
+
+ deleted = cleaner.clean_stale_gcs_buckets(self.client, now=NOW)
+
+ self.assertEqual(deleted, [])
+ self.assertFalse(value.deleted)
+
+ @mock.patch("stale_gcs_buckets_cleaner.clean_stale_gcs_buckets")
+ @mock.patch("stale_gcs_buckets_cleaner.storage.Client")
+ def test_main_defaults_to_dry_run(self, client_class, clean):
+ clean.return_value = []
+
+ with mock.patch("sys.argv", ["stale_gcs_buckets_cleaner.py"]):
+ cleaner.main()
+
+ client_class.assert_called_once_with(project="apache-beam-testing")
+ clean.assert_called_once_with(client_class.return_value, dry_run=True)
+
+ @mock.patch("stale_gcs_buckets_cleaner.clean_stale_gcs_buckets")
+ @mock.patch("stale_gcs_buckets_cleaner.storage.Client")
+ def test_main_requires_delete_flag_for_deletion(self, client_class, clean):
+ clean.return_value = []
+
+ with mock.patch(
+ "sys.argv", ["stale_gcs_buckets_cleaner.py", "--delete"]):
+ cleaner.main()
+
+ clean.assert_called_once_with(client_class.return_value, dry_run=False)
+
+
+if __name__ == "__main__":
+ unittest.main()
From 6572166744efdead363df4a832be35cfc8f40040 Mon Sep 17 00:00:00 2001
From: Maksym Tymoshyk
Date: Fri, 16 Oct 2026 17:35:07 +0300
Subject: [PATCH 011/108] [Java] Make non-extensible classes final instead of
suppressing CT_CONSTRUCTOR_THROW (#40444)
---
.../datatokenization/utils/SchemasUtils.java | 9 +--------
.../subprocess/utils/CallingSubProcessUtils.java | 9 +--------
.../subprocess/utils/ExecutableFile.java | 9 +--------
.../it/clickhouse/ClickHouseResourceManager.java | 11 ++---------
.../it/gcp/bigtable/BigtableResourceManager.java | 9 +--------
.../beam/it/splunk/SplunkResourceManager.java | 9 +--------
.../SerializablePipelineOptions.java | 9 +--------
.../runners/core/metrics/BoundedTrieData.java | 16 ++++++----------
...FlinkStreamingPortablePipelineTranslator.java | 9 +--------
.../streaming/io/UnboundedSourceWrapper.java | 10 ++--------
.../streaming/state/FlinkStateInternals.java | 9 +--------
...FlinkStreamingPortablePipelineTranslator.java | 9 +--------
.../streaming/io/UnboundedSourceWrapper.java | 10 ++--------
.../io/source/FlinkSourceSplitEnumerator.java | 9 +--------
.../source/LazyFlinkSourceSplitEnumerator.java | 9 +--------
.../stableinput/BufferingDoFnRunner.java | 9 +--------
.../KeyedBufferingElementsHandler.java | 9 +--------
.../streaming/state/FlinkStateInternals.java | 9 +--------
.../control/DefaultJobBundleFactory.java | 9 +--------
.../org/apache/beam/runners/jet/JetRunner.java | 9 +--------
.../runners/jet/processors/AssignWindowP.java | 9 +--------
.../spark/metrics/SparkBeamMetricSource.java | 9 +--------
.../beam/runners/spark/metrics/sink/CsvSink.java | 9 +--------
.../runners/spark/metrics/sink/GraphiteSink.java | 9 +--------
.../metrics/SparkBeamMetricSource.java | 9 +--------
.../metrics/sink/CodahaleCsvSink.java | 9 +--------
.../metrics/sink/CodahaleGraphiteSink.java | 9 +--------
.../SparkBatchPortablePipelineTranslator.java | 9 +--------
.../SparkDatasetPortablePipelineTranslator.java | 9 +--------
...SparkStreamingPortablePipelineTranslator.java | 9 +--------
.../org/apache/beam/sdk/transforms/Sample.java | 9 +--------
31 files changed, 39 insertions(+), 251 deletions(-)
diff --git a/examples/java/src/main/java/org/apache/beam/examples/complete/datatokenization/utils/SchemasUtils.java b/examples/java/src/main/java/org/apache/beam/examples/complete/datatokenization/utils/SchemasUtils.java
index 6388ed9fa0e7..3f64edee5c44 100644
--- a/examples/java/src/main/java/org/apache/beam/examples/complete/datatokenization/utils/SchemasUtils.java
+++ b/examples/java/src/main/java/org/apache/beam/examples/complete/datatokenization/utils/SchemasUtils.java
@@ -21,7 +21,6 @@
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import com.google.api.services.bigquery.model.TableSchema;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.io.InputStream;
import java.io.Reader;
@@ -54,13 +53,7 @@
"argument",
"return"
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SchemasUtils {
+public final class SchemasUtils {
/* Logger for class.*/
private static final Logger LOG = LoggerFactory.getLogger(SchemasUtils.class);
diff --git a/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/CallingSubProcessUtils.java b/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/CallingSubProcessUtils.java
index e04ad4c5a419..6f22bbe6855b 100644
--- a/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/CallingSubProcessUtils.java
+++ b/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/CallingSubProcessUtils.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.examples.subprocess.utils;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
@@ -89,13 +88,7 @@ private static void releaseSemaphore(String binaryName) throws IllegalStateExcep
}
/** Permit class for access to worker cpu resources. */
- @SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
- public static class Permit implements AutoCloseable {
+ public static final class Permit implements AutoCloseable {
private String binaryName;
diff --git a/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/ExecutableFile.java b/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/ExecutableFile.java
index bc409a89581c..fbff1af4afd0 100644
--- a/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/ExecutableFile.java
+++ b/examples/java/src/main/java/org/apache/beam/examples/subprocess/utils/ExecutableFile.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.examples.subprocess.utils;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import org.apache.beam.examples.subprocess.configuration.SubProcessConfiguration;
import org.apache.beam.sdk.coders.DefaultCoder;
import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
@@ -27,13 +26,7 @@
@SuppressWarnings({
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class ExecutableFile {
+public final class ExecutableFile {
String fileName;
diff --git a/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManager.java b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManager.java
index 9ff7ee92f1f6..cd7cbfa45b6f 100644
--- a/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManager.java
+++ b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManager.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.it.clickhouse;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
@@ -45,14 +44,8 @@
*
*
The class is thread-safe.
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class ClickHouseResourceManager extends TestContainerResourceManager>
- implements ResourceManager {
+public final class ClickHouseResourceManager
+ extends TestContainerResourceManager> implements ResourceManager {
private static final Logger LOG = LoggerFactory.getLogger(ClickHouseResourceManager.class);
diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigtable/BigtableResourceManager.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigtable/BigtableResourceManager.java
index 8773e2610454..eed85e8d0126 100644
--- a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigtable/BigtableResourceManager.java
+++ b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/bigtable/BigtableResourceManager.java
@@ -49,7 +49,6 @@
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
@@ -78,13 +77,7 @@
*
*
The class is thread-safe.
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class BigtableResourceManager implements ResourceManager {
+public final class BigtableResourceManager implements ResourceManager {
private static final Logger LOG = LoggerFactory.getLogger(BigtableResourceManager.class);
private static final String DEFAULT_CLUSTER_ZONE = "us-central1-b";
diff --git a/it/splunk/src/main/java/org/apache/beam/it/splunk/SplunkResourceManager.java b/it/splunk/src/main/java/org/apache/beam/it/splunk/SplunkResourceManager.java
index a4bdca3dcc31..48937fa94c06 100644
--- a/it/splunk/src/main/java/org/apache/beam/it/splunk/SplunkResourceManager.java
+++ b/it/splunk/src/main/java/org/apache/beam/it/splunk/SplunkResourceManager.java
@@ -25,7 +25,6 @@
import com.splunk.ResultsReader;
import com.splunk.ResultsReaderXml;
import com.splunk.ServiceArgs;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.time.Duration;
@@ -61,13 +60,7 @@
*
Note: The Splunk TestContainer will only run on M1 Mac's if the Docker version is >= 4.16.0
* and the "Use Rosetta for x86/amd64 emulation on Apple Silicon" setting is enabled.
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SplunkResourceManager extends TestContainerResourceManager
+public final class SplunkResourceManager extends TestContainerResourceManager
implements ResourceManager {
private static final Logger LOG = LoggerFactory.getLogger(SplunkResourceManager.class);
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/construction/SerializablePipelineOptions.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/construction/SerializablePipelineOptions.java
index 8885974c4728..f78913db3ff3 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/construction/SerializablePipelineOptions.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/construction/SerializablePipelineOptions.java
@@ -19,7 +19,6 @@
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.Serializable;
@@ -33,13 +32,7 @@
* Holds a {@link PipelineOptions} in JSON serialized form and calls {@link
* FileSystems#setDefaultPipelineOptions(PipelineOptions)} on construction or on deserialization.
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SerializablePipelineOptions implements Serializable {
+public final class SerializablePipelineOptions implements Serializable {
private static final ObjectMapper MAPPER =
new ObjectMapper()
.registerModules(ObjectMapper.findModules(ReflectHelpers.findClassLoader()));
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/BoundedTrieData.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/BoundedTrieData.java
index f99dc3e0b5de..6e55411c600e 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/BoundedTrieData.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/BoundedTrieData.java
@@ -57,13 +57,9 @@
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
@SuppressFBWarnings(
- value = {"IS2_INCONSISTENT_SYNC", "CT_CONSTRUCTOR_THROW"},
- justification =
- "Some access on purpose are left unsynchronized."
- + " Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class BoundedTrieData implements Serializable {
+ value = "IS2_INCONSISTENT_SYNC",
+ justification = "Some access on purpose are left unsynchronized")
+public final class BoundedTrieData implements Serializable {
private static final int DEFAULT_BOUND = 100; // Default maximum size of the trie
@@ -286,7 +282,7 @@ public boolean isEmpty() {
}
@Override
- public final boolean equals(@Nullable Object other) {
+ public boolean equals(@Nullable Object other) {
if (this == other) {
return true;
}
@@ -298,12 +294,12 @@ public final boolean equals(@Nullable Object other) {
}
@Override
- public final int hashCode() {
+ public int hashCode() {
return this.asTrie().hashCode();
}
@Override
- public final String toString() {
+ public String toString() {
return "BoundedTrieData(" + this.asTrie() + ")";
}
diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
index 9e24bd779de0..6ed13902d92d 100644
--- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
+++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
@@ -26,7 +26,6 @@
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.auto.service.AutoService;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
@@ -131,13 +130,7 @@
"keyfor",
"nullness"
}) // TODO(https://github.com/apache/beam/issues/20497)
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class FlinkStreamingPortablePipelineTranslator
+public final class FlinkStreamingPortablePipelineTranslator
implements FlinkPortablePipelineTranslator<
FlinkStreamingPortablePipelineTranslator.StreamingTranslationContext> {
diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
index 84790026a762..86a3f406cf0d 100644
--- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
+++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.io;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.LinkedHashMap;
@@ -68,13 +67,8 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class UnboundedSourceWrapper
+public final class UnboundedSourceWrapper<
+ OutputT, CheckpointMarkT extends UnboundedSource.CheckpointMark>
extends RichParallelSourceFunction>>
implements BeamStoppableFunction,
CheckpointListener,
diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
index a388d77c4fb0..33b098564bb5 100644
--- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
+++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.state;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.Collections;
import java.util.HashSet;
@@ -104,13 +103,7 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class FlinkStateInternals implements StateInternals {
+public final class FlinkStateInternals implements StateInternals {
private static final StateNamespace globalWindowNamespace =
StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE);
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
index 133693ec8df1..622a7c41e8d9 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
@@ -26,7 +26,6 @@
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.auto.service.AutoService;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
@@ -131,13 +130,7 @@
"keyfor",
"nullness"
}) // TODO(https://github.com/apache/beam/issues/20497)
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class FlinkStreamingPortablePipelineTranslator
+public final class FlinkStreamingPortablePipelineTranslator
implements FlinkPortablePipelineTranslator<
FlinkStreamingPortablePipelineTranslator.StreamingTranslationContext> {
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
index c5d9b8241ac3..5f1630b7a0fa 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.io;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.LinkedHashMap;
@@ -67,13 +66,8 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class UnboundedSourceWrapper
+public final class UnboundedSourceWrapper<
+ OutputT, CheckpointMarkT extends UnboundedSource.CheckpointMark>
extends RichParallelSourceFunction>>
implements BeamStoppableFunction,
CheckpointListener,
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java
index ed20c7fa5259..b3ae5b988627 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
@@ -40,13 +39,7 @@
import org.slf4j.LoggerFactory;
/** Splits a Beam source and assigns its splits to Flink source readers round-robin. */
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class FlinkSourceSplitEnumerator
+public final class FlinkSourceSplitEnumerator
implements SplitEnumerator, FlinkSourceEnumeratorState> {
private static final Logger LOG = LoggerFactory.getLogger(FlinkSourceSplitEnumerator.class);
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
index f298e6534b03..95d789178c0a 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedHashMap;
@@ -35,13 +34,7 @@
import org.slf4j.LoggerFactory;
/** Splits a bounded Beam source and assigns one split for each reader request. */
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class LazyFlinkSourceSplitEnumerator
+public final class LazyFlinkSourceSplitEnumerator
implements SplitEnumerator, FlinkSourceEnumeratorState> {
private static final Logger LOG = LoggerFactory.getLogger(LazyFlinkSourceSplitEnumerator.class);
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java
index 428ab09f5987..46aa1ee92580 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.stableinput;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -57,13 +56,7 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class BufferingDoFnRunner implements DoFnRunner {
+public final class BufferingDoFnRunner implements DoFnRunner {
public static BufferingDoFnRunner create(
DoFnRunner doFnRunner,
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/KeyedBufferingElementsHandler.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/KeyedBufferingElementsHandler.java
index 1f26c23209d2..7c2ac9f97166 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/KeyedBufferingElementsHandler.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/KeyedBufferingElementsHandler.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.stableinput;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;
@@ -32,13 +31,7 @@
@SuppressWarnings({
"rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class KeyedBufferingElementsHandler implements BufferingElementsHandler {
+public final class KeyedBufferingElementsHandler implements BufferingElementsHandler {
static KeyedBufferingElementsHandler create(
KeyedStateBackend backend, ListStateDescriptor stateDescriptor)
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
index ce8dff37b0ea..e2d1217d65bb 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.state;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.Collections;
import java.util.HashSet;
@@ -104,13 +103,7 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class FlinkStateInternals implements StateInternals {
+public final class FlinkStateInternals implements StateInternals {
private static final StateNamespace globalWindowNamespace =
StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE);
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactory.java
index eb8dca84261b..4a1adccd4e69 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactory.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.fnexecution.control;
import com.google.auto.value.AutoValue;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.IdentityHashMap;
import java.util.Map;
@@ -97,13 +96,7 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class DefaultJobBundleFactory implements JobBundleFactory {
+public final class DefaultJobBundleFactory implements JobBundleFactory {
private static final Logger LOG = LoggerFactory.getLogger(DefaultJobBundleFactory.class);
private static final IdGenerator factoryIdGenerator = IdGenerators.incrementingLongs();
diff --git a/runners/jet/src/main/java/org/apache/beam/runners/jet/JetRunner.java b/runners/jet/src/main/java/org/apache/beam/runners/jet/JetRunner.java
index 77da05d28938..97d0e6dfbc96 100644
--- a/runners/jet/src/main/java/org/apache/beam/runners/jet/JetRunner.java
+++ b/runners/jet/src/main/java/org/apache/beam/runners/jet/JetRunner.java
@@ -24,7 +24,6 @@
import com.hazelcast.jet.config.JobConfig;
import com.hazelcast.jet.core.DAG;
import com.hazelcast.map.IMap;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
@@ -48,13 +47,7 @@
import org.slf4j.LoggerFactory;
/** Jet specific implementation of Beam's {@link PipelineRunner}. */
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class JetRunner extends PipelineRunner {
+public final class JetRunner extends PipelineRunner {
private static final Logger LOG = LoggerFactory.getLogger(JetRunner.class);
diff --git a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/AssignWindowP.java b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/AssignWindowP.java
index 233266cf9fdb..b177c53ccfd6 100644
--- a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/AssignWindowP.java
+++ b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/AssignWindowP.java
@@ -21,7 +21,6 @@
import com.hazelcast.jet.core.AbstractProcessor;
import com.hazelcast.jet.core.Processor;
import com.hazelcast.jet.core.ResettableSingletonTraverser;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Collection;
import javax.annotation.Nonnull;
import org.apache.beam.runners.jet.Utils;
@@ -44,13 +43,7 @@
"nullness",
"keyfor"
}) // TODO(https://github.com/apache/beam/issues/20497)
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class AssignWindowP extends AbstractProcessor {
+public final class AssignWindowP extends AbstractProcessor {
@SuppressWarnings({"FieldCanBeLocal", "unused"})
private final String ownerId; // do not remove, useful for debugging
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/SparkBeamMetricSource.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/SparkBeamMetricSource.java
index bd93ee023396..c49094c695b9 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/SparkBeamMetricSource.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/SparkBeamMetricSource.java
@@ -18,20 +18,13 @@
package org.apache.beam.runners.spark.metrics;
import com.codahale.metrics.MetricRegistry;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import org.apache.spark.metrics.source.Source;
/**
* A Spark {@link Source} that is tailored to expose a {@link SparkBeamMetric}, wrapping an
* underlying {@link org.apache.beam.sdk.metrics.MetricResults} instance.
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SparkBeamMetricSource implements Source {
+public final class SparkBeamMetricSource implements Source {
private static final String METRIC_NAME = "Metrics";
private final String name;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/CsvSink.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/CsvSink.java
index 23635267fda4..e24e7b988caf 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/CsvSink.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/CsvSink.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.spark.metrics.sink;
import com.codahale.metrics.MetricRegistry;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Properties;
import org.apache.beam.runners.spark.metrics.WithMetricsSupport;
import org.apache.spark.SecurityManager;
@@ -37,13 +36,7 @@
* "spark.metrics.conf.*.sink.csv.unit"=seconds
* }
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class CsvSink implements Sink {
+public final class CsvSink implements Sink {
// Initialized reflectively as done by Spark's MetricsSystem
private final org.apache.spark.metrics.sink.CsvSink delegate;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/GraphiteSink.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/GraphiteSink.java
index 7ffeb7234002..c25d8c5ac291 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/GraphiteSink.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/sink/GraphiteSink.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.spark.metrics.sink;
import com.codahale.metrics.MetricRegistry;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Properties;
import org.apache.beam.runners.spark.metrics.WithMetricsSupport;
import org.apache.spark.SecurityManager;
@@ -40,13 +39,7 @@
* "spark.metrics.conf.*.sink.graphite.regex"=""
* }
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class GraphiteSink implements Sink {
+public final class GraphiteSink implements Sink {
// Initialized reflectively as done by Spark's MetricsSystem
private final org.apache.spark.metrics.sink.GraphiteSink delegate;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/SparkBeamMetricSource.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/SparkBeamMetricSource.java
index 149e132c56b3..b19be4a8ff70 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/SparkBeamMetricSource.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/SparkBeamMetricSource.java
@@ -18,20 +18,13 @@
package org.apache.beam.runners.spark.structuredstreaming.metrics;
import com.codahale.metrics.MetricRegistry;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import org.apache.spark.metrics.source.Source;
/**
* A Spark {@link Source} that is tailored to expose a {@link SparkBeamMetric}, wrapping an
* underlying {@link org.apache.beam.sdk.metrics.MetricResults} instance.
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SparkBeamMetricSource implements Source {
+public final class SparkBeamMetricSource implements Source {
private final String name;
private final MetricRegistry metricRegistry = new MetricRegistry();
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleCsvSink.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleCsvSink.java
index b46287acf6fe..905916682ea5 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleCsvSink.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleCsvSink.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.spark.structuredstreaming.metrics.sink;
import com.codahale.metrics.MetricRegistry;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Properties;
import org.apache.beam.runners.spark.structuredstreaming.metrics.WithMetricsSupport;
import org.apache.spark.SecurityManager;
@@ -37,13 +36,7 @@
* "spark.metrics.conf.*.sink.csv.unit"=seconds
* }
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class CodahaleCsvSink implements Sink {
+public final class CodahaleCsvSink implements Sink {
// Initialized reflectively as done by Spark's MetricsSystem
private final org.apache.spark.metrics.sink.CsvSink delegate;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleGraphiteSink.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleGraphiteSink.java
index 93231a758a3a..2da663ad20dd 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleGraphiteSink.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/sink/CodahaleGraphiteSink.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.spark.structuredstreaming.metrics.sink;
import com.codahale.metrics.MetricRegistry;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.Properties;
import org.apache.beam.runners.spark.structuredstreaming.metrics.WithMetricsSupport;
import org.apache.spark.SecurityManager;
@@ -40,13 +39,7 @@
* "spark.metrics.conf.*.sink.graphite.regex"=""
* }
*/
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class CodahaleGraphiteSink implements Sink {
+public final class CodahaleGraphiteSink implements Sink {
// Initialized reflectively as done by Spark's MetricsSystem
private final org.apache.spark.metrics.sink.GraphiteSink delegate;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java
index 2590fb0f8d5a..540c6d2ce0fa 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkBatchPortablePipelineTranslator.java
@@ -26,7 +26,6 @@
import static org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.instantiateCoder;
import com.google.auto.service.AutoService;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
@@ -78,13 +77,7 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SparkBatchPortablePipelineTranslator
+public final class SparkBatchPortablePipelineTranslator
implements SparkPortablePipelineTranslator {
private final ImmutableMap urnToTransformTranslator;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkDatasetPortablePipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkDatasetPortablePipelineTranslator.java
index 8aa110371390..4eaf97fa50a5 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkDatasetPortablePipelineTranslator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkDatasetPortablePipelineTranslator.java
@@ -25,7 +25,6 @@
import static org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.hasUnboundedPCollections;
import static org.apache.spark.sql.functions.col;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
@@ -82,13 +81,7 @@
"unchecked",
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SparkDatasetPortablePipelineTranslator
+public final class SparkDatasetPortablePipelineTranslator
implements SparkPortablePipelineTranslator {
private final ImmutableMap urnToTransformTranslator;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java
index dbe28ca2e23d..7bfad22d9a52 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkStreamingPortablePipelineTranslator.java
@@ -24,7 +24,6 @@
import static org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.getWindowedValueCoder;
import static org.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.getWindowingStrategy;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
@@ -73,13 +72,7 @@
"rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
-@SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
-public class SparkStreamingPortablePipelineTranslator
+public final class SparkStreamingPortablePipelineTranslator
implements SparkPortablePipelineTranslator {
private final ImmutableMap urnToTransformTranslator;
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Sample.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Sample.java
index 0d2f35a5a012..e5c18afb9e44 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Sample.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Sample.java
@@ -19,7 +19,6 @@
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
@@ -292,13 +291,7 @@ public T extractOutput(List accumulator) {
*
* @param the type of the elements
*/
- @SuppressFBWarnings(
- value = "CT_CONSTRUCTOR_THROW",
- justification =
- "Pre-existing finding, not triaged yet."
- + " Making the class final or moving the throwing code"
- + " into a static factory method may fix it.")
- public static class FixedSizedSampleFn
+ public static final class FixedSizedSampleFn
extends CombineFn<
T, Top.BoundedHeap, SerializableComparator>>, Iterable> {
private final int sampleSize;
From 1f76c9c36d642f58ca81b6b1fb757411771323cb Mon Sep 17 00:00:00 2001
From: Shunping Huang
Date: Fri, 16 Oct 2026 10:45:22 -0400
Subject: [PATCH 012/108] Fix Python unit tests broken by google-api-core
arriving via cloud-sql-python-connector 1.23.0 (#40623)
---
.../beam_PostCommit_Python_Versions.json | 2 +-
.../apache_beam/dataframe/io_it_test.py | 7 ++--
sdks/python/apache_beam/dataframe/io_test.py | 11 ++----
.../apache_beam/io/requestresponse_it_test.py | 5 +++
sdks/python/apache_beam/testing/test_utils.py | 25 ++++++++++++
.../apache_beam/testing/test_utils_test.py | 37 ++++++++++++++++++
.../enrichment_handlers/cloudsql_it_test.py | 3 ++
sdks/python/conftest.py | 38 ++++++++++++++-----
8 files changed, 106 insertions(+), 22 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Python_Versions.json b/.github/trigger_files/beam_PostCommit_Python_Versions.json
index 8b2c8c445c1f..b5704c67ef1c 100644
--- a/.github/trigger_files/beam_PostCommit_Python_Versions.json
+++ b/.github/trigger_files/beam_PostCommit_Python_Versions.json
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
- "revision": 5
+ "revision": 6
}
diff --git a/sdks/python/apache_beam/dataframe/io_it_test.py b/sdks/python/apache_beam/dataframe/io_it_test.py
index da88ffb54760..9f750e2ff58c 100644
--- a/sdks/python/apache_beam/dataframe/io_it_test.py
+++ b/sdks/python/apache_beam/dataframe/io_it_test.py
@@ -33,13 +33,12 @@
_LOGGER = logging.getLogger(__name__)
try:
- from google.api_core.exceptions import GoogleAPICallError
+ from apitools.base.py.exceptions import HttpError
except ImportError:
- GoogleAPICallError = None
+ HttpError = None
-@unittest.skipIf(
- GoogleAPICallError is None, 'GCP dependencies are not installed')
+@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class ReadUsingReadGbqTests(unittest.TestCase):
@pytest.mark.it_postcommit
def test_ReadGbq(self):
diff --git a/sdks/python/apache_beam/dataframe/io_test.py b/sdks/python/apache_beam/dataframe/io_test.py
index 051a85b379f7..d9723e1e5bb5 100644
--- a/sdks/python/apache_beam/dataframe/io_test.py
+++ b/sdks/python/apache_beam/dataframe/io_test.py
@@ -46,9 +46,9 @@
from apache_beam.testing.util import equal_to
try:
- from google.api_core.exceptions import GoogleAPICallError
+ from apitools.base.py.exceptions import HttpError
except ImportError:
- GoogleAPICallError = None
+ HttpError = None
# Get major, minor version
PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2]))
@@ -492,15 +492,10 @@ def test_double_write(self):
set(self.read_all_lines(output + 'out2.csv*')))
-@unittest.skipIf(
- GoogleAPICallError is None, 'GCP dependencies are not installed')
+@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class ReadGbqTransformTests(unittest.TestCase):
@mock.patch.object(BigQueryWrapper, 'get_table')
def test_bad_schema_public_api_direct_read(self, get_table):
- try:
- bigquery.TableFieldSchema
- except AttributeError:
- raise ValueError('Please install GCP Dependencies.')
fields = [
bigquery.TableFieldSchema(name='stn', type='DOUBLE', mode="NULLABLE"),
bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"),
diff --git a/sdks/python/apache_beam/io/requestresponse_it_test.py b/sdks/python/apache_beam/io/requestresponse_it_test.py
index 072a1dc27b9f..4e4c2b638441 100644
--- a/sdks/python/apache_beam/io/requestresponse_it_test.py
+++ b/sdks/python/apache_beam/io/requestresponse_it_test.py
@@ -30,6 +30,7 @@
from apache_beam.coders import coders
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.test_utils import require_linux_docker
# pylint: disable=ungrouped-imports
try:
@@ -206,6 +207,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
@pytest.mark.uses_testcontainer
class TestRedisCache(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls) -> None:
+ require_linux_docker()
+
def setUp(self) -> None:
self.retries = 5
self._start_container()
diff --git a/sdks/python/apache_beam/testing/test_utils.py b/sdks/python/apache_beam/testing/test_utils.py
index 425b4c6dfac4..ee953a902a0d 100644
--- a/sdks/python/apache_beam/testing/test_utils.py
+++ b/sdks/python/apache_beam/testing/test_utils.py
@@ -31,6 +31,7 @@
import struct
import sys
import tempfile
+from unittest import SkipTest
import numpy as np
@@ -186,6 +187,30 @@ def block_imports(*packages):
sys.modules.update(hidden_modules)
+def require_linux_docker():
+ """Raises ``unittest.SkipTest`` unless Docker can run Linux containers.
+
+ Intended for tests that start containers via testcontainers, e.g. from
+ ``setUpClass``. Such tests need a reachable Docker daemon running Linux
+ containers. Some environments have no daemon at all (e.g. GitHub macOS
+ runners), while others only run Windows containers (e.g. GitHub Windows
+ runners).
+ """
+ try:
+ import docker
+ client = docker.from_env()
+ try:
+ os_type = client.info().get('OSType')
+ finally:
+ client.close()
+ except Exception as e: # pylint: disable=broad-except
+ raise SkipTest(f'Docker is not available: {e}')
+ if os_type != 'linux':
+ raise SkipTest(
+ f'Docker daemon runs {os_type} containers; '
+ 'Linux containers are required.')
+
+
@retry.with_exponential_backoff(
num_retries=3, retry_filter=retry.retry_on_beam_io_error_filter)
def delete_files(file_paths):
diff --git a/sdks/python/apache_beam/testing/test_utils_test.py b/sdks/python/apache_beam/testing/test_utils_test.py
index 9c379cf35aab..3b99687c1158 100644
--- a/sdks/python/apache_beam/testing/test_utils_test.py
+++ b/sdks/python/apache_beam/testing/test_utils_test.py
@@ -126,6 +126,43 @@ def test_block_imports_only_blocks_given_packages(self):
self.assertEqual(attempted_imports, [])
+class RequireLinuxDockerTest(unittest.TestCase):
+ def _fake_docker(self, os_type=None, error=None):
+ fake_docker = mock.MagicMock()
+ client = fake_docker.from_env.return_value
+ if error:
+ client.info.side_effect = error
+ else:
+ client.info.return_value = {'OSType': os_type}
+ return fake_docker, client
+
+ def test_linux_docker_does_not_skip(self):
+ fake_docker, client = self._fake_docker(os_type='linux')
+ with mock.patch.dict(sys.modules, {'docker': fake_docker}):
+ utils.require_linux_docker()
+ client.close.assert_called_once()
+
+ def test_windows_docker_skips(self):
+ fake_docker, client = self._fake_docker(os_type='windows')
+ with mock.patch.dict(sys.modules, {'docker': fake_docker}):
+ with self.assertRaisesRegex(unittest.SkipTest, 'windows containers'):
+ utils.require_linux_docker()
+ client.close.assert_called_once()
+
+ def test_unreachable_docker_skips(self):
+ fake_docker, client = self._fake_docker(
+ error=ConnectionError('no such file'))
+ with mock.patch.dict(sys.modules, {'docker': fake_docker}):
+ with self.assertRaisesRegex(unittest.SkipTest, 'Docker is not available'):
+ utils.require_linux_docker()
+ client.close.assert_called_once()
+
+ def test_missing_docker_sdk_skips(self):
+ with utils.block_imports('docker'):
+ with self.assertRaisesRegex(unittest.SkipTest, 'Docker is not available'):
+ utils.require_linux_docker()
+
+
class LCGeneratorTest(unittest.TestCase):
Generator = utils.LCGenerator
diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql_it_test.py b/sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql_it_test.py
index 04db85a75c29..644dd56978bc 100644
--- a/sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql_it_test.py
+++ b/sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql_it_test.py
@@ -28,6 +28,7 @@
import apache_beam as beam
from apache_beam.coders import coders
from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.test_utils import require_linux_docker
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
@@ -589,6 +590,8 @@ def setUpClass(cls):
# Skip setup for the base class.
raise unittest.SkipTest("Base class - no db_adapter defined")
+ require_linux_docker()
+
# Type hint data from subclasses.
cls._db_adapter: DatabaseTypeAdapter
diff --git a/sdks/python/conftest.py b/sdks/python/conftest.py
index 92d90769fb15..0b2db2649a99 100644
--- a/sdks/python/conftest.py
+++ b/sdks/python/conftest.py
@@ -18,11 +18,11 @@
"""Pytest configuration and custom hooks."""
import gc
+import math
import os
import sys
import threading
import time
-from types import SimpleNamespace
import pytest
@@ -150,6 +150,32 @@ def enhance_mock_stability():
time.sleep(0.05)
+def _configure_testcontainers_timeouts():
+ """Applies TC_* environment settings to testcontainers' global config.
+
+ testcontainers 4.x computes its wait timeout as ``max_tries * sleep_time``
+ from ``testcontainers_config``. It reads ``TC_MAX_TRIES`` and
+ ``TC_POOLING_INTERVAL`` itself, but not ``TC_TIMEOUT`` or ``TC_SLEEP_TIME``,
+ which Beam's CI workflows set. Translate them here so that the effective
+ timeout is at least ``TC_TIMEOUT`` seconds.
+ """
+ from testcontainers.core.config import testcontainers_config
+
+ sleep_time = float(
+ os.getenv("TC_SLEEP_TIME") or os.getenv("TC_POOLING_INTERVAL") or "1")
+ timeout = float(os.getenv("TC_TIMEOUT", "120"))
+ max_tries = int(os.getenv("TC_MAX_TRIES", "0"))
+
+ testcontainers_config.sleep_time = sleep_time
+ testcontainers_config.max_tries = max(
+ max_tries, math.ceil(timeout / sleep_time))
+ print(
+ "Set testcontainers max_tries=%d, sleep_time=%s (timeout=%ss)" % (
+ testcontainers_config.max_tries,
+ testcontainers_config.sleep_time,
+ testcontainers_config.timeout))
+
+
def pytest_configure(config):
"""Saves options added in pytest_addoption for later use.
This is necessary since pytest-xdist workers do not have the
@@ -159,14 +185,8 @@ def pytest_configure(config):
# for the entire test session.
print("\n--- Applying global testcontainers timeout configuration ---")
try:
- from testcontainers.core import waiting_utils
- waiting_utils.config = SimpleNamespace(
- timeout=int(os.getenv("TC_TIMEOUT", "120")),
- max_tries=int(os.getenv("TC_MAX_TRIES", "120")),
- sleep_time=float(os.getenv("TC_SLEEP_TIME", "1")),
- )
- print("Successfully set waiting utils config")
- except ModuleNotFoundError:
+ _configure_testcontainers_timeouts()
+ except ImportError:
print("The testcontainers library is not installed.")
TestPipeline.pytest_test_pipeline_options = config.getoption(
From 084fb321e579e0c64167f59a1f949f0ed582d400 Mon Sep 17 00:00:00 2001
From: Shunping Huang
Date: Fri, 16 Oct 2026 12:24:40 -0400
Subject: [PATCH 013/108] Restore Java 11 support in setJavaVerOptions (#40634)
* Restore Java 11 support in setJavaVerOptions
the '11' branch. Workflows that still pass -PtestJavaVersion=11 now fail
while Gradle configures the projects:
Unknown Java Version 11 for setting additional java options
* Trigger test workflow
---
.github/trigger_files/beam_PreCommit_Java_IOs_Direct.json | 4 ++++
.../groovy/org/apache/beam/gradle/BeamModulePlugin.groovy | 4 +++-
2 files changed, 7 insertions(+), 1 deletion(-)
create mode 100644 .github/trigger_files/beam_PreCommit_Java_IOs_Direct.json
diff --git a/.github/trigger_files/beam_PreCommit_Java_IOs_Direct.json b/.github/trigger_files/beam_PreCommit_Java_IOs_Direct.json
new file mode 100644
index 000000000000..5abe02fc09c7
--- /dev/null
+++ b/.github/trigger_files/beam_PreCommit_Java_IOs_Direct.json
@@ -0,0 +1,4 @@
+{
+ "comment": "Modify this file in a trivial way to cause this test suite to run.",
+ "modification": 1
+}
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
index d699f9efab94..cf5979dd176b 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -1021,7 +1021,9 @@ class BeamModulePlugin implements Plugin {
}
options.fork = true
options.forkOptions.javaHome = javaVerHome as File
- if (ver == '17') {
+ if (ver == '11') {
+ // No additional options needed. Error Prone module exports are only required on JDK 16+.
+ } else if (ver == '17') {
options.forkOptions.jvmArgs += errorProneAddModuleOpts.collect { '-J' + it }
} else if (ver == '21' || ver == '25') {
options.compilerArgs += ['-Xlint:-this-escape']
From d5f6472f73a42cca0bd0bd19ffd17498fdb5f4bb Mon Sep 17 00:00:00 2001
From: tvalentyn
Date: Fri, 16 Oct 2026 10:02:41 -0700
Subject: [PATCH 014/108] Enable throttling logic in RemoteInference. (#40598)
* Enable throttler.
* Preserve while loop instead of if.
---
.../apache/beam/sdk/ml/inference/remote/RemoteInference.java | 4 ++--
.../beam/sdk/ml/inference/remote/RemoteInferenceTest.java | 3 ---
sdks/python/apache_beam/io/components/adaptive_throttler.py | 4 ++--
sdks/python/apache_beam/ml/inference/base.py | 2 +-
4 files changed, 5 insertions(+), 8 deletions(-)
diff --git a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
index a998f4d0f077..3dd4158069b0 100644
--- a/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
+++ b/sdks/java/ml/inference/remote/src/main/java/org/apache/beam/sdk/ml/inference/remote/RemoteInference.java
@@ -143,7 +143,7 @@ public Invoke withThrottleDelaySecs(int throttleDelaySecs) {
/**
* Configures the length of history to consider when setting throttling probability. Defaults to
- * a sample period of 1000ms. For more context, see {@link AdaptiveThrottler}
+ * a sample period of 120000ms. For more context, see {@link AdaptiveThrottler}
*/
public Invoke withSamplePeriodMs(long samplePeriodMs) {
checkArgument(samplePeriodMs > 0, "samplePeriodMs must be positive");
@@ -239,7 +239,7 @@ static class RemoteInferenceFninvoke()
.handler(MockThrottlingHandler.class)
.withBatchConfig(batchConfig)
- // Use large sample periods so the 1s retry delay doesn't flush the history
- .withSamplePeriodMs(60000L)
- .withSampleUpdateMs(60000L)
// Set to 1 second to minimize test wait time while still verifying throttling
.withThrottleDelaySecs(1)
.withOverloadRatio(1.1)
diff --git a/sdks/python/apache_beam/io/components/adaptive_throttler.py b/sdks/python/apache_beam/io/components/adaptive_throttler.py
index 3c22891ee8a3..06f2d83d7142 100644
--- a/sdks/python/apache_beam/io/components/adaptive_throttler.py
+++ b/sdks/python/apache_beam/io/components/adaptive_throttler.py
@@ -131,12 +131,12 @@ class ReactiveThrottler(AdaptiveThrottler):
system recovery. capture the timestamp of the attempted request, then execute
the request code. On a success, call successful_request(timestamp) to report
the success to the throttler. This flow looks like the following:
-
+
def remote_call():
throttler.throttle()
try:
- timestamp = time.time()
+ timestamp = time.time() * 1000
result = make_request()
throttler.successful_request(timestamp)
return result
diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 85a287ad7529..823a4a3fa719 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -425,7 +425,7 @@ def __init__(
throttle_delay_secs: int = 5,
retry_filter: Callable[[Exception], bool] = lambda x: True,
*,
- window_ms: int = 1 * _MILLISECOND_TO_SECOND,
+ window_ms: int = 120 * _MILLISECOND_TO_SECOND,
bucket_ms: int = 1 * _MILLISECOND_TO_SECOND,
overload_ratio: float = 2,
rate_limiter: Optional[RateLimiter] = None):
From 964d15b47a48b70d3536c33c55edbc95a4a318a3 Mon Sep 17 00:00:00 2001
From: raams95
Date: Fri, 16 Oct 2026 12:07:18 -0500
Subject: [PATCH 015/108] [Python] Document the import-time default of
PeriodicImpulse start_timestamp (#40599)
The default start_timestamp of PeriodicImpulse and MatchContinuously is
evaluated once, when the module is imported. As discussed in #40565 the
default stays as is, so document it and point to an explicit timestamp
or rebase for starting relative to pipeline execution.
Fixes #40564
Co-authored-by: JanakiRamChittineni <341781521+raams95@users.noreply.github.com>
---
sdks/python/apache_beam/io/fileio.py | 3 ++-
sdks/python/apache_beam/transforms/periodicsequence.py | 5 ++++-
2 files changed, 6 insertions(+), 2 deletions(-)
diff --git a/sdks/python/apache_beam/io/fileio.py b/sdks/python/apache_beam/io/fileio.py
index fcce83fa59ec..23d0f8c10c1a 100644
--- a/sdks/python/apache_beam/io/fileio.py
+++ b/sdks/python/apache_beam/io/fileio.py
@@ -404,7 +404,8 @@ def __init__(
file_pattern: The file path to read from.
interval: Interval at which to check for files in seconds.
has_deduplication: Whether files already read are discarded or not.
- start_timestamp: Timestamp for start file checking.
+ start_timestamp: Timestamp for start file checking. Defaults to the time
+ this module was first imported, not the time the transform is created.
stop_timestamp: Timestamp after which no more files will be checked.
match_updated_files: (When has_deduplication is set to True) whether match
file with timestamp changes.
diff --git a/sdks/python/apache_beam/transforms/periodicsequence.py b/sdks/python/apache_beam/transforms/periodicsequence.py
index 7b90fa3bb2fa..2513885e75dc 100644
--- a/sdks/python/apache_beam/transforms/periodicsequence.py
+++ b/sdks/python/apache_beam/transforms/periodicsequence.py
@@ -299,7 +299,10 @@ def __init__(
data: Optional[Sequence[Any]] = None,
rebase: RebaseMode = RebaseMode.REBASE_NONE):
'''
- :param start_timestamp: Timestamp for first element.
+ :param start_timestamp: Timestamp for first element. Defaults to the time
+ this module was first imported, not the time the transform is created
+ or the pipeline starts. Pass an explicit timestamp, or use `rebase`, to
+ start relative to pipeline execution.
:param stop_timestamp: Timestamp at or after which no elements will be
output.
:param fire_interval: Interval in seconds at which to output elements.
From a3f3c8e01d4380ce2fd56e14adeacf7c09f51313 Mon Sep 17 00:00:00 2001
From: tvalentyn
Date: Fri, 16 Oct 2026 10:32:07 -0700
Subject: [PATCH 016/108] Fix handling of cancelled tickets. (#40588)
---
.../apache_beam/ml/inference/model_manager.py | 16 +++---
.../ml/inference/model_manager_test.py | 54 +++++++++++++++++++
2 files changed, 64 insertions(+), 6 deletions(-)
diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py b/sdks/python/apache_beam/ml/inference/model_manager.py
index 69db694cc970..bc16feb6f60f 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager.py
@@ -470,12 +470,13 @@ def should_spawn_model(self, tag: str, ticket_num: int) -> bool:
return False
def _wake_next_in_queue(self):
+ # Clean up cancelled tickets at head of queue
+ while self._wait_queue and self._wait_queue[
+ 0].ticket_num in self._cancelled_tickets:
+ self._cancelled_tickets.remove(self._wait_queue[0].ticket_num)
+ heapq.heappop(self._wait_queue)
+ # The queue may be empty if all remaining tickets were cancelled.
if self._wait_queue:
- # Clean up cancelled tickets at head of queue
- while self._wait_queue and self._wait_queue[
- 0].ticket_num in self._cancelled_tickets:
- self._cancelled_tickets.remove(self._wait_queue[0].ticket_num)
- heapq.heappop(self._wait_queue)
next_inline = self._wait_queue[0]
next_inline.wake_event.set()
@@ -660,7 +661,10 @@ def _evict_to_make_space(
# TODO: Also factor in the active counts to avoid thrashing
demand_map = Counter()
for item in self._wait_queue:
- demand_map[item.tag] += 1
+ # Cancelled tickets linger in the queue until they reach the head;
+ # they no longer represent demand.
+ if item.ticket_num not in self._cancelled_tickets:
+ demand_map[item.tag] += 1
my_demand = demand_map[requesting_tag]
am_i_starving = len(self._models[requesting_tag]) == 0
diff --git a/sdks/python/apache_beam/ml/inference/model_manager_test.py b/sdks/python/apache_beam/ml/inference/model_manager_test.py
index fc84f4ba8598..d1f2b35bc1d9 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager_test.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py
@@ -30,6 +30,7 @@
from apache_beam.metrics.metricbase import MetricName
from apache_beam.ml.inference.model_manager import GPUMonitor
from apache_beam.ml.inference.model_manager import ModelManager
+ from apache_beam.ml.inference.model_manager import QueueTicket
from apache_beam.ml.inference.model_manager import ResourceEstimator
except ImportError as e:
raise unittest.SkipTest("Model Manager dependencies are not installed")
@@ -436,6 +437,32 @@ def run_inference():
est_cost = self.manager._estimator.get_estimate(model_name)
self.assertAlmostEqual(est_cost, model_cost, delta=100.0)
+ def test_wake_next_in_queue_all_cancelled(self):
+ """Draining a queue of only cancelled tickets must not raise."""
+ with self.manager._cv:
+ for ticket_num in (100, 101):
+ self.manager._wait_queue.append(QueueTicket(1, ticket_num, "m"))
+ self.manager._cancelled_tickets.add(ticket_num)
+
+ self.manager._wake_next_in_queue()
+
+ self.assertEqual(self.manager._wait_queue, [])
+ self.assertEqual(self.manager._cancelled_tickets, set())
+
+ def test_wake_next_in_queue_skips_cancelled_head(self):
+ """The first non-cancelled ticket is woken."""
+ cancelled = QueueTicket(1, 100, "m")
+ live = QueueTicket(1, 101, "m")
+ with self.manager._cv:
+ self.manager._wait_queue.extend([cancelled, live])
+ self.manager._cancelled_tickets.add(100)
+
+ self.manager._wake_next_in_queue()
+
+ self.assertEqual(self.manager._wait_queue, [live])
+ self.assertTrue(live.wake_event.is_set())
+ self.assertFalse(cancelled.wake_event.is_set())
+
class TestModelManagerEviction(unittest.TestCase):
def setUp(self):
@@ -583,6 +610,33 @@ def acquire_b():
self.manager.release_model("B", model_b)
executor.shutdown(wait=True)
+ def test_cancelled_tickets_do_not_count_as_demand(self):
+ self.manager._estimator.set_initial_estimate("A", 4000)
+ self.manager._estimator.set_initial_estimate("B", 4000)
+
+ model_a = self.manager.acquire_model("A", self.create_loader("A", 4000))
+ self.manager.release_model("A", model_a)
+ # B stays active so it is not starving and demand comparison applies.
+ model_b = self.manager.acquire_model("B", self.create_loader("B", 4000))
+
+ with self.manager._cv:
+ # One live request for B, two cancelled (timed out) requests for A.
+ self.manager._wait_queue.extend([
+ QueueTicket(1, 100, "B"),
+ QueueTicket(1, 101, "A"),
+ QueueTicket(1, 102, "A"),
+ ])
+ self.manager._cancelled_tickets.update({101, 102})
+
+ # 8000 used + 6000 requested > 12000, so A must be evicted. Without
+ # ignoring cancelled tickets, A's demand (2) >= B's (1) protects it.
+ made_space = self.manager._evict_to_make_space(
+ limit=12000, est_cost=6000, requesting_tag="B")
+
+ self.assertTrue(made_space)
+ self.assertTrue(model_a.deleted)
+ self.manager.release_model("B", model_b)
+
class TestGPUMonitor(unittest.TestCase):
def setUp(self):
From 1fc615fcd2b35b82a9b15bb3b260c989a7a64511 Mon Sep 17 00:00:00 2001
From: Derrick Williams
Date: Fri, 16 Oct 2026 14:09:36 -0400
Subject: [PATCH 017/108] [typescript] - fix typescript failures (#40635)
* [typescript] - fix typescript failures
* fix test call
---
.../src/apache_beam/internal/serialize.ts | 2 +-
.../typescript/src/apache_beam/serialization.ts | 17 +++++++++++++++--
sdks/typescript/test/primitives_test.ts | 6 ++----
3 files changed, 18 insertions(+), 7 deletions(-)
diff --git a/sdks/typescript/src/apache_beam/internal/serialize.ts b/sdks/typescript/src/apache_beam/internal/serialize.ts
index 6584e41d6af9..29f474ef4858 100644
--- a/sdks/typescript/src/apache_beam/internal/serialize.ts
+++ b/sdks/typescript/src/apache_beam/internal/serialize.ts
@@ -16,10 +16,10 @@
* limitations under the License.
*/
+import { requireForSerialization, registeredObjects } from "../serialization";
import * as serialize_closures from "serialize-closures";
import Long from "long";
-import { requireForSerialization, registeredObjects } from "../serialization";
import { packageName } from "../utils/packageJson";
const BIGINT_PREFIX = ":bigint:";
diff --git a/sdks/typescript/src/apache_beam/serialization.ts b/sdks/typescript/src/apache_beam/serialization.ts
index 9ca2177926fe..1c7ce53ca64c 100644
--- a/sdks/typescript/src/apache_beam/serialization.ts
+++ b/sdks/typescript/src/apache_beam/serialization.ts
@@ -16,10 +16,23 @@
* limitations under the License.
*/
-import * as serialize_closures from "serialize-closures";
+// Polyfill deprecated util helpers removed in Node 24+ that are still used by
+// serialize-closures.
+const nodeUtil = require("util");
+if (typeof nodeUtil.isPrimitive !== "function") {
+ nodeUtil.isPrimitive = (val: unknown): boolean =>
+ val === null || (typeof val !== "object" && typeof val !== "function");
+}
+if (typeof nodeUtil.isArray !== "function") {
+ nodeUtil.isArray = Array.isArray;
+}
+if (typeof nodeUtil.isFunction !== "function") {
+ nodeUtil.isFunction = (val: unknown): boolean => typeof val === "function";
+}
const registeredModules: string[] = [];
-export const registeredObjects: serialize_closures.BuiltinRecord[] = [];
+export const registeredObjects: import("serialize-closures").BuiltinRecord[] =
+ [];
const registeredObjectSet: Set = new Set();
/**
diff --git a/sdks/typescript/test/primitives_test.ts b/sdks/typescript/test/primitives_test.ts
index b3f35cce1fbf..b7adb28d9f03 100644
--- a/sdks/typescript/test/primitives_test.ts
+++ b/sdks/typescript/test/primitives_test.ts
@@ -308,10 +308,8 @@ export function suite(runner: beam.Runner = directRunner()) {
return v * 4;
});
- assert.deepEqual(
- p.context.getPCollectionCoder(res),
- new GeneralObjectCoder()
- );
+ const coder = p.context.getPCollectionCoder(res);
+ assert.equal(coder.constructor.name, "GeneralObjectCoder");
});
// why doesn't map need types here?
it("runs a GroupBy expansion", function () {
From a17e73a0da4710a1fee926ccd2b7abac2ea2ec08 Mon Sep 17 00:00:00 2001
From: tvalentyn
Date: Fri, 16 Oct 2026 12:17:05 -0700
Subject: [PATCH 018/108] Clarify known restrictions for model updates.
(#40600)
---
sdks/python/apache_beam/ml/inference/utils.py | 13 ++++++++-----
1 file changed, 8 insertions(+), 5 deletions(-)
diff --git a/sdks/python/apache_beam/ml/inference/utils.py b/sdks/python/apache_beam/ml/inference/utils.py
index dfbb01a29049..8cbe836e22dd 100644
--- a/sdks/python/apache_beam/ml/inference/utils.py
+++ b/sdks/python/apache_beam/ml/inference/utils.py
@@ -133,14 +133,17 @@ def __init__(
1. Any previously used filenames cannot be reused. If a file is added
or updated to a previously used filename, this transform will ignore
that update. To trigger a model update, always upload a file with
- unique name.
- 2. Initially, before the pipeline startup time, WatchFilePattern expects
+ unique name. Do not preserve old modification timestamps when copying
+ or uploading a file (for example, when rolling back to an older model),
+ as files with timestamps older than the pipeline start time are ignored.
+ 2. Do not upload more than one model update per polling ``interval``, as
+ matching multiple new files in the same poll can cause undefined
+ ordering or singleton side-input errors.
+ 3. Initially, before the pipeline startup time, WatchFilePattern expects
at least one file present that matches the file_pattern.
- 3. This transform is supported in streaming mode since
+ 4. This transform is supported in streaming mode since
MatchContinuously produces an unbounded source. Running in batch
mode can lead to undesired results or result in pipeline being stuck.
-
-
"""
self.file_pattern = file_pattern
self.interval = interval
From a74f328f26eb748ab47f2d4647c28667447cd551 Mon Sep 17 00:00:00 2001
From: Derrick Williams
Date: Fri, 16 Oct 2026 15:50:19 -0400
Subject: [PATCH 019/108] [yaml] - fix post commit yaml test (#40633)
* [yaml] - fix post commit yaml test
* trigger postcommit
---
.../beam_PostCommit_Yaml_Xlang_Direct.json | 2 +-
.../beam_PreCommit_Yaml_Xlang_Direct.json | 2 +-
.../extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml | 10 +++++-----
3 files changed, 7 insertions(+), 7 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json
index a975cd1cd104..541dc4ea8e87 100644
--- a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json
+++ b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
- "revision": 1
+ "revision": 2
}
diff --git a/.github/trigger_files/beam_PreCommit_Yaml_Xlang_Direct.json b/.github/trigger_files/beam_PreCommit_Yaml_Xlang_Direct.json
index 8c604b0a135c..840d064bdbb0 100644
--- a/.github/trigger_files/beam_PreCommit_Yaml_Xlang_Direct.json
+++ b/.github/trigger_files/beam_PreCommit_Yaml_Xlang_Direct.json
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
- "revision": 2
+ "revision": 3
}
diff --git a/sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml b/sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml
index e492e697af3a..86c1a4660613 100644
--- a/sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml
+++ b/sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml
@@ -64,7 +64,7 @@ pipelines:
key:
callable: |
def convert_to_bytes(row):
- return bytes(row.key, 'utf-8') if "key" in row._fields else None
+ return bytes(row.key, 'utf-8') if hasattr(row, "key") else None
type:
type
family_name:
@@ -72,11 +72,11 @@ pipelines:
column_qualifier:
callable: |
def convert_to_bytes(row):
- return bytes(row.column_qualifier, 'utf-8') if 'column_qualifier' in row._fields else None
+ return bytes(row.column_qualifier, 'utf-8') if hasattr(row, 'column_qualifier') else None
value:
callable: |
def convert_to_bytes(row):
- return bytes(row.value, 'utf-8') if 'value' in row._fields else None
+ return bytes(row.value, 'utf-8') if hasattr(row, 'value') else None
timestamp_micros:
timestamp_micros
- type: WriteToBigTable
@@ -102,13 +102,13 @@ pipelines:
key:
callable: |
def convert_to_string(row):
- return row.key.decode("utf-8") if "key" in row._fields else None
+ return row.key.decode("utf-8") if hasattr(row, "key") else None
family_name:
family_name
column_qualifier:
callable: |
def convert_to_string(row):
- return row.column_qualifier.decode("utf-8") if "column_qualifier" in row._fields else None
+ return row.column_qualifier.decode("utf-8") if hasattr(row, "column_qualifier") else None
cells:
callable: |
def convert_to_string(row):
From e41f0ce1d1d5526190e7e3de1820f5a7b22e05d3 Mon Sep 17 00:00:00 2001
From: Shunping Huang
Date: Fri, 16 Oct 2026 15:50:33 -0400
Subject: [PATCH 020/108] [Python] Add examples for AWS credentials from a
secret manager (#40592)
* Add examples for AWS credentials from a secret manager
Add examples under apache_beam/examples/secret_manager that set
--aws_credentials_secret_manager to read AWS credentials from Google Cloud
Secret Manager:
* aws_credentials_s3_read.py: reads an S3 file outside a pipeline.
* aws_credentials_s3_pipeline.py: reads an S3 file in a pipeline
(DirectRunner or Dataflow).
* aws_credentials_secret_read.py: reads a secret from AWS Secrets Manager
with Secret.parse_secret_option and Secret.from_json.
* Add a use case of arn for reading AWS secrets
* Reformat
* Added a requirements file for the example.
---
.../examples/secret_manager/__init__.py | 18 ++++
.../aws_credentials_requirements.txt | 21 +++++
.../aws_credentials_s3_pipeline.py | 75 ++++++++++++++++
.../secret_manager/aws_credentials_s3_read.py | 60 +++++++++++++
.../aws_credentials_secret_read.py | 87 +++++++++++++++++++
5 files changed, 261 insertions(+)
create mode 100644 sdks/python/apache_beam/examples/secret_manager/__init__.py
create mode 100644 sdks/python/apache_beam/examples/secret_manager/aws_credentials_requirements.txt
create mode 100644 sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_pipeline.py
create mode 100644 sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_read.py
create mode 100644 sdks/python/apache_beam/examples/secret_manager/aws_credentials_secret_read.py
diff --git a/sdks/python/apache_beam/examples/secret_manager/__init__.py b/sdks/python/apache_beam/examples/secret_manager/__init__.py
new file mode 100644
index 000000000000..92dc81c246f6
--- /dev/null
+++ b/sdks/python/apache_beam/examples/secret_manager/__init__.py
@@ -0,0 +1,18 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Examples for SecretManager integration in Apache Beam."""
diff --git a/sdks/python/apache_beam/examples/secret_manager/aws_credentials_requirements.txt b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_requirements.txt
new file mode 100644
index 000000000000..e0db2ef261db
--- /dev/null
+++ b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_requirements.txt
@@ -0,0 +1,21 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# Worker dependencies for aws_credentials_s3_pipeline.py: boto3 for the S3
+# filesystem, and google-cloud-secret-manager to resolve the AWS credentials.
+boto3>=1.9,<2
+google-cloud-secret-manager>=2.0,<3
diff --git a/sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_pipeline.py b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_pipeline.py
new file mode 100644
index 000000000000..535444653499
--- /dev/null
+++ b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_pipeline.py
@@ -0,0 +1,75 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""A pipeline that reads an S3 file and logs each line, using AWS credentials
+stored in Google Cloud Secret Manager.
+
+Usage (DirectRunner):
+ python -m apache_beam.examples.secret_manager.aws_credentials_s3_pipeline \
+ --input=s3://my-bucket/my-file.txt \
+ --aws_region_name=us-east-1 \
+ --aws_credentials_secret_manager=GoogleCloudSecretManager \
+ --aws_access_key_id='{"name": "my-aws-key-id", "project": "my-project"}' \
+ --aws_secret_access_key='{"name": "my-aws-secret-key", "project": "my-project"}'
+
+Usage (DataflowRunner): add the following flags to the command above.
+ --runner=DataflowRunner \
+ --project=my-project \
+ --region=us-central1 \
+ --temp_location=gs://my-gcs-bucket/temp \
+ --requirements_file=apache_beam/examples/secret_manager/aws_credentials_requirements.txt
+
+where aws_credentials_requirements.txt (in the same directory as this example)
+lists the packages the workers need: boto3 and google-cloud-secret-manager.
+
+Required permissions:
+ * The AWS IAM user that owns the access key needs ``s3:ListBucket`` on the
+ bucket (``arn:aws:s3:::my-bucket``, used to match the input path) and
+ ``s3:GetObject`` on the object (``arn:aws:s3:::my-bucket/my-file.txt``),
+ for example through the ``AmazonS3ReadOnlyAccess`` managed policy or a
+ narrower custom policy. If the object is encrypted with a customer managed
+ KMS key, it also needs ``kms:Decrypt`` on that key.
+ * The Google Cloud account that resolves the credentials needs the Secret
+ Manager Secret Accessor role (``roles/secretmanager.secretAccessor``) on
+ the two secrets that hold the access key ID and secret access key. The
+ credentials are resolved both when the pipeline is constructed and where
+ the file is read, so this is the account running the script and, on
+ Dataflow, also the worker service account.
+"""
+
+import argparse
+import logging
+
+import apache_beam as beam
+from apache_beam.options.pipeline_options import PipelineOptions
+
+
+def main(argv=None):
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input', required=True, help='S3 path to read.')
+ known_args, pipeline_args = parser.parse_known_args(argv)
+
+ with beam.Pipeline(options=PipelineOptions(pipeline_args)) as p:
+ _ = (
+ p
+ | 'Read' >> beam.io.ReadFromText(known_args.input)
+ | 'Log' >> beam.Map(logging.info))
+
+
+if __name__ == '__main__':
+ logging.getLogger().setLevel(logging.INFO)
+ main()
diff --git a/sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_read.py b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_read.py
new file mode 100644
index 000000000000..6ef1c1ba8049
--- /dev/null
+++ b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_s3_read.py
@@ -0,0 +1,60 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Reads an S3 file outside of a pipeline, using AWS credentials stored in
+Google Cloud Secret Manager.
+
+Usage:
+ python -m apache_beam.examples.secret_manager.aws_credentials_s3_read \
+ --input=s3://my-bucket/my-file.txt \
+ --aws_region_name=us-east-1 \
+ --aws_credentials_secret_manager=GoogleCloudSecretManager \
+ --aws_access_key_id='{"name": "my-aws-key-id", "project": "my-project"}' \
+ --aws_secret_access_key='{"name": "my-aws-secret-key", "project": "my-project"}'
+
+Required permissions:
+ * The AWS IAM user that owns the access key needs ``s3:GetObject`` on the
+ object (e.g. ``arn:aws:s3:::my-bucket/my-file.txt``), for example through
+ the ``AmazonS3ReadOnlyAccess`` managed policy or a narrower custom policy.
+ If the object is encrypted with a customer managed KMS key, it also needs
+ ``kms:Decrypt`` on that key.
+ * The Google Cloud account running this script (e.g. your Application
+ Default Credentials) needs the Secret Manager Secret Accessor role
+ (``roles/secretmanager.secretAccessor``) on the two secrets that hold the
+ access key ID and secret access key.
+"""
+
+import argparse
+
+from apache_beam.io.filesystems import FileSystems
+from apache_beam.options.pipeline_options import PipelineOptions
+
+
+def main(argv=None):
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input', required=True, help='S3 path to read.')
+ known_args, pipeline_args = parser.parse_known_args(argv)
+
+ # The S3 filesystem resolves the --aws_* credentials through the secret
+ # manager given by --aws_credentials_secret_manager.
+ FileSystems.set_options(PipelineOptions(pipeline_args))
+ with FileSystems.open(known_args.input) as f:
+ print(f.read().decode('utf-8'))
+
+
+if __name__ == '__main__':
+ main()
diff --git a/sdks/python/apache_beam/examples/secret_manager/aws_credentials_secret_read.py b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_secret_read.py
new file mode 100644
index 000000000000..bf3dadedeb2d
--- /dev/null
+++ b/sdks/python/apache_beam/examples/secret_manager/aws_credentials_secret_read.py
@@ -0,0 +1,87 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Reads a secret from AWS Secrets Manager outside of a pipeline, using AWS
+credentials stored in Google Cloud Secret Manager.
+
+Usage:
+ python -m apache_beam.examples.secret_manager.aws_credentials_secret_read \
+ --secret_id=my-aws-secret \
+ --secret_arn=arn:aws:secretsmanager:us-east-1::secret:my-aws-secret-AbCdEf \
+ --aws_region_name=us-east-1 \
+ --aws_credentials_secret_manager=GoogleCloudSecretManager \
+ --aws_access_key_id='{"name": "my-aws-key-id", "project": "my-project"}' \
+ --aws_secret_access_key='{"name": "my-aws-secret-key", "project": "my-project"}'
+
+Required permissions:
+ * The AWS IAM user that owns the access key needs
+ ``secretsmanager:GetSecretValue`` on the secret (e.g.
+ ``arn:aws:secretsmanager:us-east-1::secret:my-aws-secret-*``),
+ for example through the ``AWSSecretsManagerClientReadOnlyAccess`` managed
+ policy or a narrower custom policy. If the secret is encrypted with a
+ customer managed KMS key instead of the default ``aws/secretsmanager`` key,
+ it also needs ``kms:Decrypt`` on that key.
+ * The Google Cloud account running this script (e.g. your Application
+ Default Credentials) needs the Secret Manager Secret Accessor role
+ (``roles/secretmanager.secretAccessor``) on the two secrets that hold the
+ access key ID and secret access key.
+"""
+
+import argparse
+
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.utils.secret import Secret
+
+
+def main(argv=None):
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ '--secret_id',
+ required=True,
+ help='Name of the secret in AWS Secrets Manager.')
+ parser.add_argument(
+ '--secret_arn', help='Optional ARN of the secret in AWS Secrets Manager.')
+ known_args, pipeline_args = parser.parse_known_args(argv)
+ options = PipelineOptions(pipeline_args)
+
+ # The AWS Secrets Manager client uses the --aws_* options, with credentials
+ # resolved through --aws_credentials_secret_manager.
+
+ # Option 1: a 'type:;:' secret string.
+ secret = Secret.parse_secret_option(
+ f'type:awssecret;name:{known_args.secret_id}', options)
+ print(f'Retrieved secret of length {len(secret.get_str())}')
+
+ # Option 2: a JSON spec and a secret manager name.
+ secret = Secret.from_json(
+ f'{{"name": "{known_args.secret_id}"}}', 'AwsSecretsManager', options)
+ print(f'Retrieved secret of length {len(secret.get_str())}')
+
+ # Options 3 and 4: the same as options 1 and 2, but identifying the secret by
+ # its ARN.
+ if known_args.secret_arn:
+ secret = Secret.parse_secret_option(
+ f'type:awssecret;arn:{known_args.secret_arn}', options)
+ print(f'Retrieved secret of length {len(secret.get_str())}')
+
+ secret = Secret.from_json(
+ f'{{"arn": "{known_args.secret_arn}"}}', 'AwsSecretsManager', options)
+ print(f'Retrieved secret of length {len(secret.get_str())}')
+
+
+if __name__ == '__main__':
+ main()
From 8eda31b281f292854e83e3f6a3d84218007eacfb Mon Sep 17 00:00:00 2001
From: Divyansh Yadav
Date: Sat, 17 Oct 2026 02:47:40 +0530
Subject: [PATCH 021/108] [Python] Fix max_models_per_worker_hint not being
enforced in KeyedModelHandler (#40500)
* [Python] Fix max_models_per_worker_hint not being enforced in KeyedModelHandler
* Guard increment_max_models with a lock on the model handler manager
---
CHANGES.md | 1 +
sdks/python/apache_beam/ml/inference/base.py | 31 +++++++++++-----
.../apache_beam/ml/inference/base_test.py | 35 +++++++++++++++++++
3 files changed, 59 insertions(+), 8 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 537c3a32447c..d9ca6f0b83dd 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -104,6 +104,7 @@
* (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)).
* (Go) Wrap last dataWriter/timerWriter send errors on Close ([#40405](https://github.com/apache/beam/issues/40405)).
* (Go) Fixed a deadlock recreating a data channel while holding the channel lock ([#40414](https://github.com/apache/beam/issues/40414)).
+* (Python) Fixed `KeyedModelHandler` not enforcing `max_models_per_worker_hint` when several copies of the model handler run in the same SDK process ([#40468](https://github.com/apache/beam/issues/40468)).
* (Java) Fixed the declared schema of the error output of the Kafka write SchemaTransform, which wrapped the error schema a second time and did not match the rows it emits ([#39760](https://github.com/apache/beam/issues/39760)).
* (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed in recent releases, which broke builds of Go modules depending on a current `genproto` version ([#40018](https://github.com/apache/beam/issues/40018)).
* (Python) `VertexAIFeatureStoreEnrichmentHandler` now logs a warning instead of failing when the caller lacks permission to read the feature store metadata ([#40442](https://github.com/apache/beam/issues/40442)).
diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 823a4a3fa719..ce014e4b04f9 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -577,6 +577,11 @@ def __init__(self, mh_map: dict[str, ModelHandler]):
model.
"""
self._max_models = None
+ # Process ids that have already contributed to _max_models
+ self._incremented_process_ids: set[int] = set()
+ # Guards _max_models and _incremented_process_ids, since the shared
+ # manager serves calls from several processes on separate threads.
+ self._max_models_lock = threading.Lock()
# Map keys to model handlers
self._mh_map: dict[str, ModelHandler] = mh_map
# Map keys to the last updated model path for that key
@@ -631,17 +636,26 @@ def load(self, key: str) -> _ModelLoadStats:
return _ModelLoadStats(
tag, end_time - start_time, memory_after - memory_before)
- def increment_max_models(self, increment: int):
+ def increment_max_models(
+ self, increment: int, process_id: Optional[int] = None):
"""
Increments the number of models that this instance of a
_ModelHandlerManager is able to hold. If it is never called,
no limit is imposed.
Args:
increment: the amount by which we are incrementing the number of models.
+ process_id: the id of the process requesting the increment. If set,
+ only the first increment from each process is applied, since many
+ copies of the model handler can share this manager within a process.
"""
- if self._max_models is None:
- self._max_models = 0
- self._max_models += increment
+ with self._max_models_lock:
+ if process_id is not None:
+ if process_id in self._incremented_process_ids:
+ return
+ self._incremented_process_ids.add(process_id)
+ if self._max_models is None:
+ self._max_models = 0
+ self._max_models += increment
def update_model_handler(self, key: str, model_path: str, previous_key: str):
"""
@@ -867,11 +881,12 @@ def run_inference(
self._unkeyed.run_inference(unkeyed_batch, model, inference_args))
# The first time a MultiProcessShared ModelManager is used for inference
- # from this process, we should increment its max model count
+ # from this process, we should increment its max model count. The manager
+ # may live in another process and be shared by several deserialized copies
+ # of this handler, so it dedupes increments by process id.
if self._max_models_per_worker_hint is not None:
- lock = threading.Lock()
- if lock.acquire(blocking=False):
- model.increment_max_models(self._max_models_per_worker_hint)
+ model.increment_max_models(
+ self._max_models_per_worker_hint, process_id=os.getpid())
self._max_models_per_worker_hint = None
batch_by_key = defaultdict(list)
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py
index 568bdf7ea129..64f5c3b1525d 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -1978,6 +1978,41 @@ def test_model_handler_manager_evicts_models_after_being_incremented(self):
mh3.load_model, tag=tag3).acquire()
self.assertEqual(8, model3.predict(10))
+ def test_model_handler_manager_increments_once_per_process(self):
+ mhs = {
+ 'key1': FakeModelHandler(state=1),
+ 'key2': FakeModelHandler(state=2),
+ 'key3': FakeModelHandler(state=3)
+ }
+ mm = base._ModelHandlerManager(mh_map=mhs)
+ mm.increment_max_models(1, process_id=1)
+ mm.increment_max_models(1, process_id=1)
+ mm.load('key1')
+ mm.load('key2')
+ mm.load('key3')
+ self.assertEqual(['key3'], list(mm._tag_map.keys()))
+
+ mm.increment_max_models(1, process_id=2)
+ mm.load('key1')
+ self.assertEqual(['key3', 'key1'], list(mm._tag_map.keys()))
+
+ def test_keyed_model_handler_max_models_hint_with_handler_copies(self):
+ mhs = [
+ base.KeyModelMapping([k],
+ FakeModelHandler(
+ state=i, multi_process_shared=True))
+ for i, k in enumerate(['a', 'b', 'c'])
+ ]
+ keyed_mh = base.KeyedModelHandler(mhs, max_models_per_worker_hint=1)
+ mm = keyed_mh.load_model()
+ # Each DoFn instance in a process deserializes its own copy of the model
+ # handler, but they all share the same _ModelHandlerManager.
+ for _ in range(5):
+ mh_copy = pickle.loads(pickle.dumps(keyed_mh))
+ mh_copy.override_metrics('test_namespace')
+ list(mh_copy.run_inference([('a', 1), ('b', 2), ('c', 3)], mm))
+ self.assertEqual(1, len(mm._tag_map))
+
def test_run_inference_loads_different_models(self):
mh1 = FakeModelHandler(incrementing=True, min_batch_size=3)
with TestPipeline() as pipeline:
From 2be52540457e3887080a67641240b243ea6dad45 Mon Sep 17 00:00:00 2001
From: pablo rodriguez defino
Date: Fri, 16 Oct 2026 17:00:26 -0700
Subject: [PATCH 022/108] [Prism] Finalize bundles of DoFns fused after another
transform (#40640)
* [Prism] Finalize bundles of DoFns fused after another transform
Problem:
In finalizeStage (sdks/go/pkg/beam/runners/prism/internal/preprocess.go), Prism sets
stage.finalize only while it inspects the consumers of the stage's input PCollections.
The loop skips every PCollection that a transform of the same stage produces:
for pid, plinks := range stageFacts.PcolConsumers {
if _, ok := stageFacts.PcolProducers[pid]; ok {
continue
}
...
if pardo.GetRequestsFinalization() {
stg.finalize = true
}
So a ParDo that sets requests_finalization but is fused after another transform (a Map
after Create, the expand step of a Reshuffle, any element-wise chain) never marks its
stage. The SDK harness registers the callbacks and returns requires_finalization, but Prism
never sends FinalizeBundleRequest, and the callbacks never run. This affects every SDK. The
existing TestParDoBundleFinalizer cases apply the finalizing DoFn directly to Impulse, so the
DoFn always consumes the stage input and the tests pass.
Fix:
Set stage.finalize from every ParDo-like transform of the stage (ParDo and the splittable
DoFn stage URNs), not only from the consumers of the stage inputs. Statefulness, the
on-window-expiration timer and side inputs still come from the input consumers. Add a
TestParDoBundleFinalizer case with the finalizing DoFn fused after another ParDo.
Reproduction (Python SDK on PrismRunner):
```python
import os
import tempfile
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
class MarkFinalization(beam.DoFn):
"""Registers a bundle finalization callback that creates the file `path`."""
def __init__(self, path):
self._path = path
def process(self, element, bundle_finalizer=beam.DoFn.BundleFinalizerParam):
path = self._path
bundle_finalizer.register(lambda: open(path, "w").close())
yield element
if __name__ == "__main__":
markers = tempfile.mkdtemp()
at_input = os.path.join(markers, "at_stage_input")
fused = os.path.join(markers, "after_fused_map")
options = PipelineOptions(["--runner=PrismRunner", "--environment_type=LOOPBACK"])
with beam.Pipeline(options=options) as p:
impulse = p | beam.Impulse()
# Consumes the input of its stage: Prism finalizes the bundle.
_ = impulse | "AtStageInput" >> beam.ParDo(MarkFinalization(at_input))
# Fused after a Map in the same stage: Prism never finalizes the bundle.
_ = (
impulse
| "Map" >> beam.Map(lambda x: x)
| "AfterFusedMap" >> beam.ParDo(MarkFinalization(fused))
)
print("finalized:", sorted(os.listdir(markers)))
assert os.path.exists(at_input), "the DoFn at the stage input was not finalized"
assert os.path.exists(fused), "the DoFn fused after Map was not finalized"
```
Before the fix: `finalized: ['at_stage_input']`, and the second assertion fails.
After the fix: `finalized: ['after_fused_map', 'at_stage_input']`.
* [Prism] Test the finalization flag of a stage, and drop an unreachable error
Move the check of requests_finalization into requestsFinalization. finalizeStage runs
computeFacts first, which already rejects a ParDo payload that does not decode, so the
error return of the earlier loop could not run.
Add TestFinalizeStage_RequestsFinalization: the flag stays false when no transform asks for
it, and is set when the transform at the stage input, a ParDo fused after another
transform, or a fused splittable DoFn process stage asks for it.
---
.../beam/runners/prism/internal/preprocess.go | 26 +++++++-
.../runners/prism/internal/preprocess_test.go | 65 +++++++++++++++++++
sdks/go/test/integration/primitives/pardo.go | 15 +++++
.../test/integration/primitives/pardo_test.go | 5 ++
4 files changed, 108 insertions(+), 3 deletions(-)
diff --git a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go
index 3311bcced9f4..8f6000276a7f 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/preprocess.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/preprocess.go
@@ -446,6 +446,15 @@ func finalizeStage(stg *stage, comps *pipepb.Components, pipelineFacts *fusionFa
transformSet[tid] = true
}
+ // Any ParDo in the stage can register bundle finalization callbacks, not only the
+ // transforms that consume the stage's inputs, so every one of them must be checked.
+ for _, tid := range stg.transforms {
+ if requestsFinalization(comps.GetTransforms()[tid]) {
+ stg.finalize = true
+ break
+ }
+ }
+
// Now we can see which consumers (inputs) aren't covered by the producers (outputs).
mainInputs := map[string]string{}
var sideInputs []engine.LinkID
@@ -468,9 +477,6 @@ func finalizeStage(stg *stage, comps *pipepb.Components, pipelineFacts *fusionFa
if err := (proto.UnmarshalOptions{}).Unmarshal(t.GetSpec().GetPayload(), pardo); err != nil {
return fmt.Errorf("unable to decode ParDoPayload for %v", link.Transform)
}
- if pardo.GetRequestsFinalization() {
- stg.finalize = true
- }
if len(pardo.GetTimerFamilySpecs())+len(pardo.GetStateSpecs())+len(pardo.GetOnWindowExpirationTimerFamilySpec()) > 0 {
stg.stateful = true
}
@@ -535,6 +541,20 @@ func finalizeStage(stg *stage, comps *pipepb.Components, pipelineFacts *fusionFa
return nil
}
+// requestsFinalization reports whether t is a ParDo, or a stage of an expanded splittable
+// DoFn, whose payload requests bundle finalization. computeFacts has already rejected
+// undecodable ParDo payloads, so a payload that does not decode reports false.
+func requestsFinalization(t *pipepb.PTransform) bool {
+ switch t.GetSpec().GetUrn() {
+ case urns.TransformParDo, urns.TransformProcessSizedElements, urns.TransformPairWithRestriction, urns.TransformSplitAndSize, urns.TransformTruncate:
+ pardo := &pipepb.ParDoPayload{}
+ err := (proto.UnmarshalOptions{}).Unmarshal(t.GetSpec().GetPayload(), pardo)
+ return err == nil && pardo.GetRequestsFinalization()
+ default:
+ return false
+ }
+}
+
// greedyFusion produces a pipeline as tightly fused as possible.
//
// Fusion is a critical optimization for performance of pipeline execution.
diff --git a/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go b/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go
index 8baaa8538690..f4064e814b39 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/preprocess_test.go
@@ -789,3 +789,68 @@ func TestComputeFacts(t *testing.T) {
})
}
}
+
+func TestFinalizeStage_RequestsFinalization(t *testing.T) {
+ pardo := func(urn string, finalize bool) *pipepb.FunctionSpec {
+ return &pipepb.FunctionSpec{
+ Urn: urn,
+ Payload: protox.MustEncode(&pipepb.ParDoPayload{RequestsFinalization: finalize}),
+ }
+ }
+ // first consumes the stage input n0; second is fused after it and consumes n1.
+ stageOf := func(first, second *pipepb.FunctionSpec) *pipepb.Components {
+ return &pipepb.Components{
+ Transforms: map[string]*pipepb.PTransform{
+ "first": {
+ Spec: first,
+ Inputs: map[string]string{"i0": "n0"},
+ Outputs: map[string]string{"o0": "n1"},
+ },
+ "second": {
+ Spec: second,
+ Inputs: map[string]string{"i0": "n1"},
+ Outputs: map[string]string{"o0": "n2"},
+ },
+ },
+ }
+ }
+
+ tests := []struct {
+ name string
+ comps *pipepb.Components
+ wantFinalize bool
+ }{
+ {
+ name: "no_transform_requests_finalization",
+ comps: stageOf(&pipepb.FunctionSpec{Urn: urns.TransformAssignWindows}, pardo(urns.TransformParDo, false)),
+ }, {
+ name: "transform_at_stage_input",
+ comps: stageOf(pardo(urns.TransformParDo, true), pardo(urns.TransformParDo, false)),
+ wantFinalize: true,
+ }, {
+ name: "pardo_fused_after_another_transform",
+ comps: stageOf(pardo(urns.TransformParDo, false), pardo(urns.TransformParDo, true)),
+ wantFinalize: true,
+ }, {
+ name: "sdf_process_fused_after_another_transform",
+ comps: stageOf(pardo(urns.TransformParDo, false), pardo(urns.TransformProcessSizedElements, true)),
+ wantFinalize: true,
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ topological := []string{"first", "second"}
+ facts, err := computeFacts(topological, test.comps)
+ if err != nil {
+ t.Fatalf("computeFacts() = %v", err)
+ }
+ stg := &stage{transforms: topological}
+ if err := finalizeStage(stg, test.comps, facts); err != nil {
+ t.Fatalf("finalizeStage() = %v", err)
+ }
+ if stg.finalize != test.wantFinalize {
+ t.Errorf("stage.finalize = %v, want %v", stg.finalize, test.wantFinalize)
+ }
+ })
+ }
+}
diff --git a/sdks/go/test/integration/primitives/pardo.go b/sdks/go/test/integration/primitives/pardo.go
index dc59d8f67b80..a799f142745a 100644
--- a/sdks/go/test/integration/primitives/pardo.go
+++ b/sdks/go/test/integration/primitives/pardo.go
@@ -34,6 +34,7 @@ func init() {
register.Function3x2(asymJoinFn)
register.Function5x0(splitByName)
register.Function2x0(emitPipelineOptions)
+ register.Function1x1(passBytesFn)
register.DoFn2x0[beam.BundleFinalization, []byte]((*processElemBundleFinalizer)(nil))
register.DoFn2x0[beam.BundleFinalization, []byte]((*finalizerInFinishBundle)(nil))
register.DoFn2x0[beam.BundleFinalization, []byte]((*finalizerInAll)(nil))
@@ -213,6 +214,20 @@ func ParDoProcessElementBundleFinalizer(s beam.Scope) {
beam.ParDo0(s, &processElemBundleFinalizer{}, imp)
}
+// ParDoBundleFinalizerAfterFusedParDo creates a beam.Pipeline with a beam.ParDo0 that processes
+// a DoFn with a beam.BundleFinalization in its ProcessElement method, after another ParDo.
+// The runner fuses both ParDos into one stage, so the DoFn that requests finalization does not
+// consume the input of the stage.
+func ParDoBundleFinalizerAfterFusedParDo(s beam.Scope) {
+ imp := beam.Impulse(s)
+ passed := beam.ParDo(s, passBytesFn, imp)
+ beam.ParDo0(s, &processElemBundleFinalizer{}, passed)
+}
+
+func passBytesFn(b []byte) []byte {
+ return b
+}
+
type processElemBundleFinalizer struct {
}
diff --git a/sdks/go/test/integration/primitives/pardo_test.go b/sdks/go/test/integration/primitives/pardo_test.go
index aa6cb3de2008..1c169272b049 100644
--- a/sdks/go/test/integration/primitives/pardo_test.go
+++ b/sdks/go/test/integration/primitives/pardo_test.go
@@ -64,6 +64,11 @@ func TestParDoBundleFinalizer(t *testing.T) {
pipelineFn: ParDoProcessElementBundleFinalizer,
want: BundleFinalizerProcess,
},
+ {
+ name: "InProcessElementAfterFusedParDo",
+ pipelineFn: ParDoBundleFinalizerAfterFusedParDo,
+ want: BundleFinalizerProcess,
+ },
{
name: "InFinishBundle",
pipelineFn: ParDoFinishBundleFinalizer,
From 8ef717637e9718e86e06ddcd14179b2b559c5df1 Mon Sep 17 00:00:00 2001
From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com>
Date: Sun, 18 Oct 2026 22:55:20 -0400
Subject: [PATCH 023/108] Bump minimum cryptography bounds beyond known severe
CVEs (#40591)
---
sdks/python/setup.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 24e9440361ca..08c5964898de 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -416,7 +416,7 @@ def get_portability_package_data():
},
ext_modules=extensions,
install_requires=[
- 'cryptography>=39.0.0',
+ 'cryptography>=50.0.0',
'fastavro>=0.23.6,<2',
'fasteners>=0.3,<1.0',
'grpcio>=1.33.1,<2,!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,!=1.66.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*', # pylint: disable=line-too-long
From 40df757ea39be36ac71beebf849a5a7997ae0ff9 Mon Sep 17 00:00:00 2001
From: Shahar Epstein <60007259+shahar1@users.noreply.github.com>
Date: Mon, 19 Oct 2026 06:09:18 +0300
Subject: [PATCH 024/108] Support attributing BigQuery quota and billing to a
custom GCP project (Python, Java, Go) (#39203)
* [Python] Support a custom quota project for BigQuery IO
Add a quota_project_id parameter to ReadFromBigQuery and a
--quota_project_id pipeline option (GoogleCloudOptions) that attribute
BigQuery API quota and billing to a project other than the one the data
resides in.
The quota project is sent as the x-goog-user-project header on the
apitools BigQuery client, and through client_options on the
google-cloud-bigquery Client (streaming inserts) and the BigQuery
Storage Read client. The shared pipeline credentials are left untouched
instead of being copied with with_quota_project(), which for some
credential types drops the cached token. The header is only added to
requests to the BigQuery API hosts: the credentials refresh their token
through the same Http, and the token endpoint, metadata server and IAM
Credentials API must not receive it.
It is honored by ReadFromBigQuery (EXPORT and DIRECT_READ, including
temporary dataset cleanup), ReadAllFromBigQuery, WriteToBigQuery with
FILE_LOADS or STREAMING_INSERTS, and ReadBigQueryChangeHistory. The
value is resolved once at expand() time and passed to each DoFn so it
reaches separately deserialized SDF copies. STORAGE_WRITE_API is
Java-backed and not covered yet.
Generated-By: Claude Opus 5.5
* [Java] Add bigQueryQuotaProjectId pipeline option to BigQueryIO
Add BigQueryOptions.bigQueryQuotaProjectId to attribute BigQuery API
quota and billing to a project other than the one the data resides in.
For the REST client, the x-goog-user-project header is set by an
execute interceptor, so it is re-applied on every attempt, including
after HttpCredentialsAdapter re-initializes headers on a 401 refresh.
The Storage Read and Write clients set it through setQuotaProjectId.
Generated-By: Claude Opus 5.5
* [Go] Add WithQuotaProject option to bigqueryio
Add a WithQuotaProject query option that attributes BigQuery API quota
and billing to a project other than the one the data resides in, via
option.WithQuotaProject on the BigQuery client.
Read now accepts query options too. Only WithQuotaProject applies to it;
UseStandardSQL and WithQueryParameters panic, because Read generates a
legacy SQL statement without parameters.
Generated-By: Claude Opus 5.5
* Add CHANGES.md entry for BigQuery quota project support
Generated-By: Claude Opus 5.5
---
CHANGES.md | 2 +-
sdks/go/pkg/beam/io/bigqueryio/bigquery.go | 54 ++-
.../pkg/beam/io/bigqueryio/bigquery_test.go | 125 ++++++
.../sdk/io/gcp/bigquery/BigQueryOptions.java | 10 +
.../io/gcp/bigquery/BigQueryServicesImpl.java | 23 +-
.../bigquery/BigQueryServicesImplTest.java | 54 +++
sdks/python/apache_beam/io/gcp/bigquery.py | 121 ++++--
.../io/gcp/bigquery_change_history.py | 42 ++-
.../io/gcp/bigquery_change_history_test.py | 75 ++++
.../apache_beam/io/gcp/bigquery_file_loads.py | 46 ++-
.../io/gcp/bigquery_file_loads_test.py | 47 +++
.../io/gcp/bigquery_read_internal.py | 11 +-
.../apache_beam/io/gcp/bigquery_test.py | 355 ++++++++++++++++++
.../apache_beam/io/gcp/bigquery_tools.py | 162 +++++++-
.../apache_beam/io/gcp/bigquery_tools_test.py | 232 +++++++++++-
.../apache_beam/options/pipeline_options.py | 11 +
16 files changed, 1282 insertions(+), 88 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index d9ca6f0b83dd..c10b7262209c 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -66,7 +66,7 @@
* Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
* (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)).
-
+* Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--quota_project_id` pipeline option, honored by the Python BigQuery IO transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet `STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)).
## New Features / Improvements
diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
index bbe951969b92..98818ccf1af1 100644
--- a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
+++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go
@@ -34,6 +34,7 @@ import (
bq "google.golang.org/api/bigquery/v2"
"google.golang.org/api/googleapi"
"google.golang.org/api/iterator"
+ "google.golang.org/api/option"
)
// writeSizeLimit is the maximum number of rows allowed by BQ in a write.
@@ -88,14 +89,22 @@ func NewQualifiedTableName(s string) (QualifiedTableName, error) {
// Read reads all rows from the given table. The table must have a schema
// compatible with the given type, t, and Read returns a PCollection. If the
// table has more rows than t, then Read is implicitly a projection.
-func Read(s beam.Scope, project, table string, t reflect.Type) beam.PCollection {
+// Of the options, only WithQuotaProject applies to Read; UseStandardSQL and
+// WithQueryParameters panic, because the generated statement is legacy SQL and
+// takes no parameters.
+func Read(s beam.Scope, project, table string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection {
mustParseTable(table)
s = s.Scope("bigquery.Read")
+ qo := applyQueryOptions(options)
+ if qo.UseStandardSQL || len(qo.parameters) > 0 {
+ panic("bigqueryio.Read: UseStandardSQL and WithQueryParameters are only supported by Query")
+ }
+
stmt := constructSelectStatement(t, bigQueryTag, table)
- return query(s, project, stmt, t)
+ return queryWithOptions(s, project, stmt, t, qo)
}
func constructSelectStatement(t reflect.Type, tagKey string, table string) string {
@@ -114,6 +123,9 @@ func constructSelectStatement(t reflect.Type, tagKey string, table string) strin
type QueryOptions struct {
// UseStandardSQL enables BigQuery's Standard SQL dialect when executing a query.
UseStandardSQL bool
+ // QuotaProject is the GCP project ID used for quota and billing attribution
+ // of the BigQuery API calls, if different from the project the data resides in.
+ QuotaProject string
// Parameters are the query parameters for parameterized queries.
// In the current implementation, user-defines types are not supported in Value field.
// Use *bigquery.QueryParameterValue to build STRUCT/ARRAY parameters
@@ -129,6 +141,17 @@ func UseStandardSQL() func(qo *QueryOptions) error {
}
}
+// WithQuotaProject sets the GCP project ID used for quota and billing
+// attribution of the BigQuery API calls, if different from the project the
+// data resides in. The credentials used must have the
+// serviceusage.services.use permission on that project.
+func WithQuotaProject(project string) func(qo *QueryOptions) error {
+ return func(qo *QueryOptions) error {
+ qo.QuotaProject = project
+ return nil
+ }
+}
+
// WithQueryParameters sets the query parameters for parameterized queries.
// See QueryOptions.parameters for the list of supported Value types.
func WithQueryParameters(params ...bigquery.QueryParameter) func(qo *QueryOptions) error {
@@ -138,6 +161,15 @@ func WithQueryParameters(params ...bigquery.QueryParameter) func(qo *QueryOption
}
}
+// clientOptions returns the BigQuery client options implied by qo.
+func clientOptions(qo QueryOptions) []option.ClientOption {
+ var opts []option.ClientOption
+ if qo.QuotaProject != "" {
+ opts = append(opts, option.WithQuotaProject(qo.QuotaProject))
+ }
+ return opts
+}
+
// Query executes a query. The output must have a schema compatible with the given
// type, t. It returns a PCollection.
func Query(s beam.Scope, project, q string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection {
@@ -145,15 +177,22 @@ func Query(s beam.Scope, project, q string, t reflect.Type, options ...func(*Que
return query(s, project, q, t, options...)
}
-func query(s beam.Scope, project, query string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection {
- mustInferSchema(t)
-
+func applyQueryOptions(options []func(*QueryOptions) error) QueryOptions {
queryOptions := QueryOptions{}
for _, opt := range options {
if err := opt(&queryOptions); err != nil {
panic(err)
}
}
+ return queryOptions
+}
+
+func query(s beam.Scope, project, query string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection {
+ return queryWithOptions(s, project, query, t, applyQueryOptions(options))
+}
+
+func queryWithOptions(s beam.Scope, project, query string, t reflect.Type, queryOptions QueryOptions) beam.PCollection {
+ mustInferSchema(t)
imp := beam.Impulse(s)
queryParameters, err := encodeQueryParameters(queryOptions.parameters)
@@ -181,8 +220,11 @@ type queryFn struct {
Options QueryOptions `json:"options"`
}
+// newClient creates the BigQuery client used by queryFn. Tests replace it.
+var newClient = bigquery.NewClient
+
func (f *queryFn) ProcessElement(ctx context.Context, _ []byte, emit func(beam.X)) error {
- client, err := bigquery.NewClient(ctx, f.Project)
+ client, err := newClient(ctx, f.Project, clientOptions(f.Options)...)
if err != nil {
return err
}
diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
index 9f83eeb847f7..4cf0c46b9f62 100644
--- a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
+++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go
@@ -16,11 +16,16 @@
package bigqueryio
import (
+ "context"
+ "errors"
+ "fmt"
"reflect"
+ "strings"
"testing"
"cloud.google.com/go/bigquery"
"github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "google.golang.org/api/option"
)
func TestNewQualifiedTableName(t *testing.T) {
@@ -149,6 +154,88 @@ func Test_mustInferSchema(t *testing.T) {
}
}
+func TestWithQuotaProject(t *testing.T) {
+ qo := QueryOptions{}
+ if err := WithQuotaProject("quota-project")(&qo); err != nil {
+ t.Fatalf("WithQuotaProject() returned err: %v", err)
+ }
+ if got, want := qo.QuotaProject, "quota-project"; got != want {
+ t.Errorf("qo.QuotaProject = %q, want %q", got, want)
+ }
+}
+
+func TestClientOptions(t *testing.T) {
+ if got := clientOptions(QueryOptions{}); len(got) != 0 {
+ t.Errorf("clientOptions(no quota) = %v, want none", got)
+ }
+ got := clientOptions(QueryOptions{QuotaProject: "quota-project"})
+ if want := []option.ClientOption{option.WithQuotaProject("quota-project")}; !reflect.DeepEqual(got, want) {
+ t.Errorf("clientOptions(with quota) = %v, want %v", got, want)
+ }
+}
+
+func TestQueryFn_passesQuotaProjectToClient(t *testing.T) {
+ var gotProject string
+ var gotOpts []option.ClientOption
+ errStop := errors.New("stop")
+ defer func(orig func(context.Context, string, ...option.ClientOption) (*bigquery.Client, error)) {
+ newClient = orig
+ }(newClient)
+ newClient = func(_ context.Context, project string, opts ...option.ClientOption) (*bigquery.Client, error) {
+ gotProject, gotOpts = project, opts
+ return nil, errStop
+ }
+
+ fn := &queryFn{Project: "project", Options: QueryOptions{QuotaProject: "quota-project"}}
+ if err := fn.ProcessElement(context.Background(), nil, nil); !errors.Is(err, errStop) {
+ t.Fatalf("ProcessElement() err = %v, want %v", err, errStop)
+ }
+ if gotProject != "project" {
+ t.Errorf("client project = %q, want %q", gotProject, "project")
+ }
+ if want := []option.ClientOption{option.WithQuotaProject("quota-project")}; !reflect.DeepEqual(gotOpts, want) {
+ t.Errorf("client options = %v, want %v", gotOpts, want)
+ }
+}
+
+// queryFnOf returns the queryFn in pipeline p.
+func queryFnOf(t *testing.T, p *beam.Pipeline) *queryFn {
+ t.Helper()
+ edges, _, err := p.Build()
+ if err != nil {
+ t.Fatalf("Build() failed: %v", err)
+ }
+ for _, e := range edges {
+ if e.DoFn != nil {
+ if fn, ok := e.DoFn.Recv.(*queryFn); ok {
+ return fn
+ }
+ }
+ }
+ t.Fatal("no queryFn in pipeline")
+ return nil
+}
+
+func TestReadAndQuery_passQuotaProjectToQueryFn(t *testing.T) {
+ type quotaSchema struct {
+ Name bigquery.NullString `bigquery:"name"`
+ }
+ beam.RegisterType(reflect.TypeOf((*quotaSchema)(nil)).Elem())
+ typ := reflect.TypeOf(quotaSchema{})
+
+ p, s := beam.NewPipelineWithRoot()
+ Read(s, "project", "project:dataset.table", typ, WithQuotaProject("quota-project"))
+ if got := queryFnOf(t, p).Options.QuotaProject; got != "quota-project" {
+ t.Errorf("Read: queryFn.Options.QuotaProject = %q, want %q", got, "quota-project")
+ }
+
+ p, s = beam.NewPipelineWithRoot()
+ Query(s, "project", "SELECT 1", typ, WithQuotaProject("quota-project"))
+ if got := queryFnOf(t, p).Options.QuotaProject; got != "quota-project" {
+ t.Errorf("Query: queryFn.Options.QuotaProject = %q, want %q", got, "quota-project")
+ }
+}
+
func TestWithQueryParameters(t *testing.T) {
t.Run("WithQueryParameters sets parameters correctly", func(t *testing.T) {
params := []bigquery.QueryParameter{
@@ -183,3 +270,41 @@ func TestWithQueryParameters(t *testing.T) {
}
})
}
+
+func TestRead_queryOnlyOptions(t *testing.T) {
+ type readSchema struct {
+ Name bigquery.NullString `bigquery:"name"`
+ }
+ beam.RegisterType(reflect.TypeOf((*readSchema)(nil)).Elem())
+ typ := reflect.TypeOf(readSchema{})
+
+ mustPanic := func(t *testing.T, opt func(*QueryOptions) error) {
+ t.Helper()
+ defer func() {
+ r := recover()
+ if r == nil {
+ t.Errorf("Read() with query-only option does not panic")
+ } else if !strings.Contains(fmt.Sprint(r), "only supported by Query") {
+ t.Errorf("Read() panicked with unexpected value: %v", r)
+ }
+ }()
+ _, s := beam.NewPipelineWithRoot()
+ Read(s, "project", "project:dataset.table", typ, opt)
+ }
+
+ t.Run("UseStandardSQL panics", func(t *testing.T) {
+ mustPanic(t, UseStandardSQL())
+ })
+ t.Run("WithQueryParameters panics", func(t *testing.T) {
+ mustPanic(t, WithQueryParameters(bigquery.QueryParameter{Name: "p", Value: 1}))
+ })
+ t.Run("WithQuotaProject does not panic", func(t *testing.T) {
+ defer func() {
+ if r := recover(); r != nil {
+ t.Errorf("Read() with WithQuotaProject panicked: %v", r)
+ }
+ }()
+ _, s := beam.NewPipelineWithRoot()
+ Read(s, "project", "project:dataset.table", typ, WithQuotaProject("quota-project"))
+ })
+}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
index 1826f4a1fd21..64382d64bb51 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
@@ -197,6 +197,16 @@ public interface BigQueryOptions
void setBigQueryProject(String value);
+ @Description(
+ "GCP project ID used for quota and billing attribution of BigQuery API requests "
+ + "(sets the X-Goog-User-Project header), if different from the project the data "
+ + "resides in. If unspecified, the project associated with the credentials is used. "
+ + "The credentials used must have the serviceusage.services.use permission on this "
+ + "project.")
+ String getBigQueryQuotaProjectId();
+
+ void setBigQueryQuotaProjectId(String value);
+
@Description("Maximum (best effort) size of a single append to the storage API.")
@Default.Integer(2 * 1024 * 1024)
Integer getStorageApiAppendThresholdBytes();
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
index 7a924c6aa021..a921be3700ec 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
@@ -1800,7 +1800,8 @@ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) throws Inte
}
/** Returns a BigQuery client builder using the specified {@link BigQueryOptions}. */
- private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
+ @VisibleForTesting
+ static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
// Do not log 404. It clutters the output and is possibly even required by the
// caller.
RetryHttpRequestInitializer httpRequestInitializer =
@@ -1818,6 +1819,18 @@ private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
initBuilder.add(new LatencyRecordingHttpRequestInitializer(API_METRIC_LABEL));
initBuilder.add(httpRequestInitializer);
+ // Set the quota project as a request header instead of deriving credentials: a derived
+ // credential inherits the shared credential's cached request metadata, which lacks
+ // x-goog-user-project until the next token refresh. Applied by an execute interceptor so it
+ // is re-applied on every attempt; HttpCredentialsAdapter re-initializes the request headers
+ // from the credential after a 401 refresh, which would otherwise override it.
+ @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId();
+ if (!Strings.isNullOrEmpty(quotaProjectId)) {
+ initBuilder.add(
+ request ->
+ request.setInterceptor(
+ r -> r.getHeaders().set("x-goog-user-project", quotaProjectId)));
+ }
HttpRequestInitializer chainInitializer =
new ChainingHttpRequestInitializer(
Iterables.toArray(initBuilder.build(), HttpRequestInitializer.class));
@@ -1848,6 +1861,10 @@ private static BigQueryWriteClient newBigQueryWriteClient(BigQueryOptions option
if (!Strings.isNullOrEmpty(endpoint)) {
builder.setEndpoint(trimSchemaIfNecessary(endpoint));
}
+ @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId();
+ if (!Strings.isNullOrEmpty(quotaProjectId)) {
+ builder.setQuotaProjectId(quotaProjectId);
+ }
return BigQueryWriteClient.create(
builder
.setCredentialsProvider(() -> options.as(GcpOptions.class).getGcpCredential())
@@ -1973,6 +1990,10 @@ public void onRetryAttempt(Status status, Metadata metadata) {
if (!Strings.isNullOrEmpty(endpoint)) {
settingsBuilder.setEndpoint(trimSchemaIfNecessary(endpoint));
}
+ @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId();
+ if (!Strings.isNullOrEmpty(quotaProjectId)) {
+ settingsBuilder.setQuotaProjectId(quotaProjectId);
+ }
UnaryCallSettings.Builder createReadSessionSettings =
settingsBuilder.getStubSettingsBuilder().createReadSessionSettings();
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
index 90b2e1fb8723..642d35b43905 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
@@ -41,6 +41,8 @@
import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo;
import com.google.api.client.googleapis.json.GoogleJsonErrorContainer;
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpResponseException;
import com.google.api.client.http.LowLevelHttpResponse;
import com.google.api.client.json.GenericJson;
@@ -2311,6 +2313,58 @@ public RetryInfo parseBytes(byte[] serialized) {
assertEquals(123456, (long) container.getCounter(metricName).getCumulative());
}
+ @Test
+ public void testQuotaProjectIdOverrides() throws IOException {
+ BigQueryOptions options = PipelineOptionsFactory.create().as(BigQueryOptions.class);
+ options.setBigQueryQuotaProjectId("my-quota-project");
+
+ assertEquals(
+ "my-quota-project",
+ new BigQueryServicesImpl.StorageClientImpl(options)
+ .getClient()
+ .getSettings()
+ .getQuotaProjectId());
+ assertEquals(
+ "my-quota-project",
+ new BigQueryServicesImpl.WriteStreamServiceImpl(options)
+ .getClient()
+ .getSettings()
+ .getQuotaProjectId());
+ }
+
+ @Test
+ public void testQuotaProjectIdSetsRequestHeader() throws IOException {
+ BigQueryOptions options = PipelineOptionsFactory.create().as(BigQueryOptions.class);
+ options.setGcpCredential(null);
+ GenericUrl url = new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects");
+
+ HttpRequest request =
+ BigQueryServicesImpl.newBigQueryClient(options)
+ .build()
+ .getRequestFactory()
+ .buildGetRequest(url);
+ request.getInterceptor().intercept(request);
+ assertNull(request.getHeaders().getFirstHeaderStringValue("x-goog-user-project"));
+
+ options.setBigQueryQuotaProjectId("my-quota-project");
+ request =
+ BigQueryServicesImpl.newBigQueryClient(options)
+ .build()
+ .getRequestFactory()
+ .buildGetRequest(url);
+ // The interceptor runs before every attempt, including retries.
+ request.getInterceptor().intercept(request);
+ assertEquals(
+ "my-quota-project", request.getHeaders().getFirstHeaderStringValue("x-goog-user-project"));
+
+ // Simulate HttpCredentialsAdapter re-initializing headers from a credential that carries
+ // its own quota project after a 401 refresh: the configured value must win on the retry.
+ request.getHeaders().set("x-goog-user-project", "credential-project");
+ request.getInterceptor().intercept(request);
+ assertEquals(
+ "my-quota-project", request.getHeaders().getFirstHeaderStringValue("x-goog-user-project"));
+ }
+
@Test
public void testEndpointOverrides() throws IOException {
BigQueryOptions options = PipelineOptionsFactory.create().as(BigQueryOptions.class);
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py
index 2a52b31edf7f..f3d64bfcfa6c 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -720,7 +720,8 @@ def __init__(
step_name=None,
unique_id=None,
temp_dataset=None,
- query_priority=BigQueryQueryPriority.BATCH):
+ query_priority=BigQueryQueryPriority.BATCH,
+ quota_project_id=None):
if table is not None and query is not None:
raise ValueError(
'Both a BigQuery table and a query were specified.'
@@ -754,6 +755,7 @@ def __init__(
self.use_json_exports = use_json_exports
self.temp_dataset = temp_dataset
self.query_priority = query_priority
+ self.quota_project_id = quota_project_id
self._job_name = job_name or 'BQ_EXPORT_JOB'
self._step_name = step_name
self._source_uuid = unique_id
@@ -773,12 +775,15 @@ def display_data(self):
'use_legacy_sql': self.use_legacy_sql,
'bigquery_job_labels': json.dumps(self.bigquery_job_labels),
'export_file_format': export_format,
+ 'quota_project_id': DisplayDataItem(
+ self._get_quota_project_id(),
+ label='BigQuery Quota Project').drop_if_none(),
'launchesBigQueryJobs': DisplayDataItem(
True, label="This Dataflow job launches bigquery jobs."),
}
def estimate_size(self):
- bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(self.options)
+ bq = self._create_bq_wrapper()
if self.table_reference is not None:
table_ref = self.table_reference
if (isinstance(self.table_reference, vp.ValueProvider) and
@@ -844,6 +849,23 @@ def _get_project(self):
project = self.project
return project
+ def _get_quota_project_id(self):
+ """Returns the explicit quota project ID, else the pipeline option."""
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ self.options)
+ return self.quota_project_id or quota_project_id
+
+ def _create_bq_wrapper(self, **kwargs):
+ """Creates a BigQueryWrapper for the API calls made by this source.
+
+ Every BigQuery client this source uses is built here, so that quota
+ attribution cannot be forgotten at an individual call site.
+ """
+ return bigquery_tools.BigQueryWrapper(
+ pipeline_options=self.options,
+ quota_project_id=self.quota_project_id,
+ **kwargs)
+
def _create_source(self, path, coder):
if not self.use_json_exports:
return create_avro_source(path, validate=self.validate)
@@ -858,10 +880,9 @@ def _create_source(self, path, coder):
def split(self, desired_bundle_size, start_position=None, stop_position=None):
if self.export_result is None:
- bq = bigquery_tools.BigQueryWrapper(
+ bq = self._create_bq_wrapper(
temp_dataset_id=(
- self.temp_dataset.datasetId if self.temp_dataset else None),
- client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options))
+ self.temp_dataset.datasetId if self.temp_dataset else None))
if self.query is not None:
self._setup_temporary_dataset(bq)
@@ -1051,7 +1072,8 @@ def __init__(
temp_dataset: Optional[DatasetReference] = None,
temp_table: Optional[TableReference] = None,
use_native_datetime: Optional[bool] = False,
- timeout: Optional[float] = None):
+ timeout: Optional[float] = None,
+ quota_project_id: Optional[str] = None):
if table is not None and query is not None:
raise ValueError(
@@ -1090,6 +1112,7 @@ def __init__(
self._job_name = job_name or 'BQ_DIRECT_READ_JOB'
self._step_name = step_name
self._source_uuid = unique_id
+ self.quota_project_id = quota_project_id
def _get_project(self):
"""Returns the project that queries and exports will be billed to."""
@@ -1101,6 +1124,23 @@ def _get_project(self):
return project
return self.project
+ def _get_quota_project_id(self):
+ """Returns the explicit quota project ID, else the pipeline option."""
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ self.pipeline_options)
+ return self.quota_project_id or quota_project_id
+
+ def _create_bq_wrapper(self, **kwargs):
+ """Creates a BigQueryWrapper for the API calls made by this source.
+
+ Every BigQuery client this source uses is built here, so that quota
+ attribution cannot be forgotten at an individual call site.
+ """
+ return bigquery_tools.BigQueryWrapper(
+ pipeline_options=self.pipeline_options,
+ quota_project_id=self.quota_project_id,
+ **kwargs)
+
def _get_parent_project(self):
"""Returns the project that will be billed."""
if self.temp_table:
@@ -1188,14 +1228,16 @@ def display_data(self):
'use_native_datetime': self.use_native_datetime,
'selected_fields': str(self.selected_fields),
'row_restriction': str(self.row_restriction),
+ 'quota_project_id': DisplayDataItem(
+ self._get_quota_project_id(),
+ label='BigQuery Quota Project').drop_if_none(),
'launchesBigQueryJobs': DisplayDataItem(
True, label="This Dataflow job launches bigquery jobs."),
}
def estimate_size(self):
# Returns the pre-filtering size of the (temporary) table being read.
- bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(
- self.pipeline_options)
+ bq = self._create_bq_wrapper()
if self.table_reference is not None:
table_ref = self.table_reference
if (isinstance(self.table_reference, vp.ValueProvider) and
@@ -1242,10 +1284,8 @@ def estimate_size(self):
def split(self, desired_bundle_size, start_position=None, stop_position=None):
if self.split_result is None:
- bq = bigquery_tools.BigQueryWrapper(
- temp_table_ref=(self.temp_table if self.temp_table else None),
- client=bigquery_tools.BigQueryWrapper._bigquery_client(
- self.pipeline_options))
+ bq = self._create_bq_wrapper(
+ temp_table_ref=(self.temp_table if self.temp_table else None))
if self.query is not None:
self._setup_temporary_dataset(bq)
@@ -1278,7 +1318,8 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None):
if self.row_restriction is not None:
requested_session.read_options.row_restriction = self.row_restriction
- storage_client = bq_storage.BigQueryReadClient()
+ storage_client = bigquery_tools.create_bq_storage_client(
+ self._get_quota_project_id())
stream_count = self._get_stream_count(bq, desired_bundle_size)
parent = 'projects/{}'.format(self.table_reference.projectId)
@@ -1302,8 +1343,10 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None):
self.split_result = [
_CustomBigQueryStorageStreamSource(
- stream.name, self.use_native_datetime, self.timeout)
- for stream in read_session.streams
+ stream.name,
+ self.use_native_datetime,
+ self.timeout,
+ self._get_quota_project_id()) for stream in read_session.streams
]
for source in self.split_result:
@@ -1337,10 +1380,12 @@ def __init__(
self,
read_stream_name: str,
use_native_datetime: Optional[bool] = True,
- timeout: Optional[float] = None):
+ timeout: Optional[float] = None,
+ quota_project_id: Optional[str] = None):
self.read_stream_name = read_stream_name
self.use_native_datetime = use_native_datetime
self.timeout = timeout
+ self.quota_project_id = quota_project_id
def display_data(self):
return {
@@ -1363,7 +1408,10 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None):
return SourceBundle(
weight=1.0,
source=_CustomBigQueryStorageStreamSource(
- self.read_stream_name, self.use_native_datetime),
+ self.read_stream_name,
+ self.use_native_datetime,
+ self.timeout,
+ self.quota_project_id),
start_position=None,
stop_position=None)
@@ -1399,7 +1447,8 @@ def retry_delay_callback(delay):
def read_arrow(self):
- storage_client = bq_storage.BigQueryReadClient()
+ storage_client = bigquery_tools.create_bq_storage_client(
+ self.quota_project_id)
read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback}
if self.timeout is not None:
read_rows_kwargs['timeout'] = self.timeout
@@ -1418,7 +1467,8 @@ def read_arrow(self):
yield py_row
def read_avro(self):
- storage_client = bq_storage.BigQueryReadClient()
+ storage_client = bigquery_tools.create_bq_storage_client(
+ self.quota_project_id)
read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback}
if self.timeout is not None:
read_rows_kwargs['timeout'] = self.timeout
@@ -1499,7 +1549,8 @@ def __init__(
with_batched_input=False,
ignore_unknown_columns=False,
max_retries=MAX_INSERT_RETRIES,
- max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE):
+ max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE,
+ quota_project_id=None):
"""Initialize a WriteToBigQuery transform.
Args:
@@ -1552,9 +1603,12 @@ def __init__(
backoffs (effectively retry forever).
max_insert_payload_size: The maximum byte size for a BigQuery legacy
streaming insert payload.
+ quota_project_id: Optional GCP project ID to use for quota and billing of
+ BigQuery API requests.
"""
self.schema = schema
self.test_client = test_client
+ self.quota_project_id = quota_project_id
self.create_disposition = create_disposition
self.write_disposition = write_disposition
if write_disposition in (BigQueryDisposition.WRITE_EMPTY,
@@ -1633,7 +1687,7 @@ def start_bundle(self):
if not self.bigquery_wrapper:
self.bigquery_wrapper = bigquery_tools.BigQueryWrapper(
- client=self.test_client)
+ client=self.test_client, quota_project_id=self.quota_project_id)
(
bigquery_tools.BigQueryWrapper.HISTOGRAM_METRIC_LOGGER.
@@ -1958,6 +2012,8 @@ def process(self, element):
yield (key, (value, insert_id))
def expand(self, input):
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ input.pipeline.options)
bigquery_write_fn = BigQueryWriteFn(
schema=self.schema,
batch_size=self.batch_size,
@@ -1971,7 +2027,8 @@ def expand(self, input):
ignore_unknown_columns=self.ignore_unknown_columns,
with_batched_input=self.with_auto_sharding,
max_retries=self._max_retries,
- max_insert_payload_size=self._max_insert_payload_size)
+ max_insert_payload_size=self._max_insert_payload_size,
+ quota_project_id=quota_project_id)
def _add_random_shard(element):
key = element[0]
@@ -3043,6 +3100,12 @@ class ReadFromBigQuery(PTransform):
PCollection with a schema and yielding Beam Rows via the option
`BEAM_ROW`. For more information on schemas, see
https://beam.apache.org/documentation/programming-guide/#what-is-a-schema)
+ quota_project_id (str): The GCP project ID to use for quota and billing
+ of BigQuery API requests issued by this transform, if different from
+ the project the data resides in. Falls back to the
+ ``--quota_project_id`` pipeline option if not set. The credentials
+ used must have the ``serviceusage.services.use`` permission on that
+ project.
query_output_schema: Required when output_type is 'BEAM_ROW' and a query
is specified. A BigQuery schema describing the query result columns,
since the schema cannot be auto-derived from an existing table when
@@ -3143,10 +3206,12 @@ def _expand_output_type(self, output_pcollection):
'%s: table must be of type string'
'; got a callable instead' % self.__class__.__name__)
return output_pcollection | bigquery_schema_tools.convert_to_usertype(
- bigquery_tools.BigQueryWrapper().get_table(
- project_id=table_details.projectId,
- dataset_id=table_details.datasetId,
- table_id=table_details.tableId).schema,
+ bigquery_tools.BigQueryWrapper(
+ pipeline_options=output_pcollection.pipeline.options,
+ quota_project_id=self._kwargs.get('quota_project_id')).get_table(
+ project_id=table_details.projectId,
+ dataset_id=table_details.datasetId,
+ table_id=table_details.tableId).schema,
self._kwargs.get('selected_fields', None))
else:
raise ValueError(
@@ -3228,7 +3293,9 @@ def _get_pipeline_details(unused_elm):
bigquery_dataset_labels=self.bigquery_dataset_labels,
*self._args,
**self._kwargs))
- | _PassThroughThenCleanupTempDatasets(project_to_cleanup_pcoll))
+ | _PassThroughThenCleanupTempDatasets(
+ project_to_cleanup_pcoll,
+ quota_project_id=self._kwargs.get('quota_project_id')))
class ReadFromBigQueryRequest:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
index 1e90c2974778..9f836190d9ca 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
@@ -416,7 +416,8 @@ def __init__(
start_time: Timestamp,
stop_time: Timestamp,
poll_interval: Duration,
- location: Optional[str] = None) -> None:
+ location: Optional[str] = None,
+ quota_project_id: Optional[str] = None) -> None:
self._table = table
self._project = project
self._change_function = change_function
@@ -425,9 +426,11 @@ def __init__(
self._stop_time = stop_time
self._poll_interval = poll_interval
self._location = location
+ self._quota_project_id = quota_project_id
def setup(self) -> None:
- self._bq_wrapper = bigquery_tools.BigQueryWrapper()
+ self._bq_wrapper = bigquery_tools.BigQueryWrapper(
+ quota_project_id=self._quota_project_id)
if self._location is None:
table_ref = bigquery_tools.parse_table_reference(
self._table, project=self._project)
@@ -570,7 +573,8 @@ def __init__(
change_type_column: str = 'change_type',
change_timestamp_column: str = 'change_timestamp',
columns: Optional[list[str]] = None,
- row_filter: Optional[str] = None) -> None:
+ row_filter: Optional[str] = None,
+ quota_project_id: Optional[str] = None) -> None:
self._table = table
self._project = project
self._change_function = change_function
@@ -580,9 +584,11 @@ def __init__(
self._change_timestamp_column = change_timestamp_column
self._columns = columns
self._row_filter = row_filter
+ self._quota_project_id = quota_project_id
def setup(self) -> None:
- self._bq_wrapper = bigquery_tools.BigQueryWrapper()
+ self._bq_wrapper = bigquery_tools.BigQueryWrapper(
+ quota_project_id=self._quota_project_id)
if self._location is None:
table_ref = bigquery_tools.parse_table_reference(
self._table, project=self._project)
@@ -711,11 +717,13 @@ def __init__(
batch_arrow_read: bool = True,
change_timestamp_column: str = 'change_timestamp',
max_split_rounds: int = 1,
- emit_raw_batches: bool = False) -> None:
+ emit_raw_batches: bool = False,
+ quota_project_id: Optional[str] = None) -> None:
self._batch_arrow_read = batch_arrow_read
self._change_timestamp_column = change_timestamp_column
self._max_split_rounds = max_split_rounds
self._emit_raw_batches = emit_raw_batches
+ self._quota_project_id = quota_project_id
self._storage_client = None
def _ensure_client(self) -> None:
@@ -727,7 +735,8 @@ def _ensure_client(self) -> None:
"""
if self._storage_client is None:
_LOGGER.info('[Read] creating BigQueryReadClient')
- self._storage_client = bq_storage.BigQueryReadClient()
+ self._storage_client = bigquery_tools.create_bq_storage_client(
+ self._quota_project_id)
def setup(self) -> None:
self._ensure_client()
@@ -1072,9 +1081,13 @@ class _CleanupTempTablesFn(beam.DoFn):
STREAMS_READ = beam.transforms.userstate.CombiningValueStateSpec(
'streams_read', sum)
+ def __init__(self, quota_project_id: Optional[str] = None) -> None:
+ self._quota_project_id = quota_project_id
+
def setup(self) -> None:
_LOGGER.info('[Cleanup] setup: creating BigQueryWrapper')
- self._bq_wrapper = bigquery_tools.BigQueryWrapper()
+ self._bq_wrapper = bigquery_tools.BigQueryWrapper(
+ quota_project_id=self._quota_project_id)
def process(
self,
@@ -1247,6 +1260,9 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection:
'project must be specified either in ReadBigQueryChangeHistory '
'or in pipeline options (--project)')
+ quota_project_id = bigquery_tools._quota_project_id_from_options(
+ pbegin.pipeline.options)
+
start_time = Timestamp(self._start_time or time.time())
stop_time = (
Timestamp(self._stop_time)
@@ -1290,7 +1306,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection:
start_time=start_time,
stop_time=stop_time,
poll_interval=poll_interval,
- location=self._location)))
+ location=self._location,
+ quota_project_id=quota_project_id)))
# CommitQueryResults: Reshuffle commits _QueryResult (temp table ref)
# so that if the Read SDF retries, it re-reads the existing temp table
@@ -1311,7 +1328,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection:
change_type_column=self._change_type_column,
change_timestamp_column=self._change_timestamp_column,
columns=self._columns,
- row_filter=self._row_filter))
+ row_filter=self._row_filter,
+ quota_project_id=quota_project_id))
| 'CommitQueryResults' >> beam.Reshuffle())
emit_raw = self._reshuffle_decompress
@@ -1321,7 +1339,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection:
batch_arrow_read=self._batch_arrow_read,
change_timestamp_column=self._change_timestamp_column,
max_split_rounds=self._max_split_rounds,
- emit_raw_batches=emit_raw))
+ emit_raw_batches=emit_raw,
+ quota_project_id=quota_project_id))
if emit_raw:
read_sdf = read_sdf.with_output_types(tuple[bytes, bytes])
else:
@@ -1334,7 +1353,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection:
_ = (
read_outputs[_CLEANUP_TAG]
- | 'CleanupTempTables' >> beam.ParDo(_CleanupTempTablesFn()))
+ | 'CleanupTempTables' >> beam.ParDo(
+ _CleanupTempTablesFn(quota_project_id=quota_project_id)))
if emit_raw:
# Reshuffle raw Arrow batches for fan-out, then decompress and
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
index 11b467f26d49..e1b1eb69ae38 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
@@ -23,12 +23,18 @@
- ReadBigQueryChangeHistory validation
"""
+import contextlib
import datetime
import unittest
+from unittest import mock
+from apache_beam.io.gcp import bigquery_change_history as bqch
+from apache_beam.io.gcp import bigquery_tools
from apache_beam.io.gcp.bigquery_change_history import ReadBigQueryChangeHistory
from apache_beam.io.gcp.bigquery_change_history import build_changes_query
from apache_beam.io.gcp.bigquery_change_history import compute_ranges
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp
@@ -205,5 +211,74 @@ def test_default_buffer(self):
self.assertEqual(t._buffer_sec, 10)
+@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
+class QuotaProjectTest(unittest.TestCase):
+ """Tests that the quota project reaches every BigQuery client."""
+ _QUOTA = 'my-billing-project'
+
+ def _assert_wrapper_quota(self, fn):
+ with mock.patch.object(bigquery_tools, 'BigQueryWrapper') as mock_wrapper:
+ fn.setup()
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'], self._QUOTA)
+
+ def test_poll_fn_setup_passes_quota_project(self):
+ fn = bqch._PollChangeHistoryFn(
+ table='proj:dataset.table',
+ project='p',
+ change_function='APPENDS',
+ buffer=Duration(seconds=10),
+ start_time=Timestamp(0),
+ stop_time=Timestamp(100),
+ poll_interval=Duration(seconds=60),
+ quota_project_id=self._QUOTA)
+ self._assert_wrapper_quota(fn)
+
+ def test_execute_query_fn_setup_passes_quota_project(self):
+ fn = bqch._ExecuteQueryFn(
+ table='proj:dataset.table',
+ project='p',
+ change_function='APPENDS',
+ temp_dataset='tmp',
+ location='US',
+ quota_project_id=self._QUOTA)
+ self._assert_wrapper_quota(fn)
+
+ def test_cleanup_fn_setup_passes_quota_project(self):
+ self._assert_wrapper_quota(bqch._CleanupTempTablesFn(self._QUOTA))
+
+ def test_read_sdf_creates_storage_client_with_quota_project(self):
+ fn = bqch._ReadStorageStreamsSDF(quota_project_id=self._QUOTA)
+ with mock.patch.object(bigquery_tools,
+ 'create_bq_storage_client') as mock_create:
+ fn.setup()
+ mock_create.assert_called_once_with(self._QUOTA)
+
+ def test_expand_resolves_quota_project_from_options(self):
+ options = PipelineOptions(
+ ['--quota_project_id=' + self._QUOTA, '--no_pipeline_type_check'])
+ classes = [
+ '_PollChangeHistoryFn',
+ '_ExecuteQueryFn',
+ '_ReadStorageStreamsSDF',
+ '_CleanupTempTablesFn'
+ ]
+ p = TestPipeline(options=options)
+ with contextlib.ExitStack() as stack:
+ mocks = [
+ stack.enter_context(
+ mock.patch.object(
+ getattr(bqch, name),
+ '__init__',
+ autospec=True,
+ side_effect=getattr(bqch, name).__init__)) for name in classes
+ ]
+ _ = p | ReadBigQueryChangeHistory(
+ table='p:d.t', project='p', start_time=0, temp_dataset='tmp')
+ for name, m in zip(classes, mocks):
+ self.assertEqual(
+ m.call_args.kwargs['quota_project_id'], self._QUOTA, name)
+
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
index dabe80c5c27e..79940d109f07 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
@@ -383,16 +383,19 @@ def __init__(
test_client=None,
additional_bq_parameters=None,
step_name=None,
- load_job_project_id=None):
+ load_job_project_id=None,
+ quota_project_id=None):
self.project = project
self._test_client = test_client
+ self._quota_project_id = quota_project_id
self._write_disposition = write_disposition
self._additional_bq_parameters = additional_bq_parameters or {}
self._step_name = step_name
self._load_job_project_id = load_job_project_id
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self._test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self._test_client, quota_project_id=self._quota_project_id)
self._bq_io_metadata = create_bigquery_io_metadata(self._step_name)
self.pending_jobs = []
@@ -525,11 +528,13 @@ def __init__(
write_disposition=None,
test_client=None,
step_name=None,
- load_job_project_id=None):
+ load_job_project_id=None,
+ quota_project_id=None):
self.project = project
self.create_disposition = create_disposition
self.write_disposition = write_disposition
self.test_client = test_client
+ self.quota_project_id = quota_project_id
self._observed_tables = set()
self.bq_io_metadata = None
self._step_name = step_name
@@ -545,7 +550,8 @@ def setup(self):
self._observed_tables = set()
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client, quota_project_id=self.quota_project_id)
if not self.bq_io_metadata:
self.bq_io_metadata = create_bigquery_io_metadata(self._step_name)
self.pending_jobs = []
@@ -677,10 +683,12 @@ def __init__(
additional_bq_parameters=None,
source_format=None,
step_name=None,
- load_job_project_id=None):
+ load_job_project_id=None,
+ quota_project_id=None):
self.schema = schema
self.project = project
self.test_client = test_client
+ self.quota_project_id = quota_project_id
self.temporary_tables = temporary_tables
self.additional_bq_parameters = additional_bq_parameters or {}
self.source_format = source_format
@@ -709,7 +717,8 @@ def display_data(self):
return result
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client, quota_project_id=self.quota_project_id)
if not self.bq_io_metadata:
self.bq_io_metadata = create_bigquery_io_metadata(self._step_name)
self.pending_jobs = []
@@ -928,11 +937,13 @@ def process(self, element):
class DeleteTablesFn(beam.DoFn):
- def __init__(self, test_client=None):
+ def __init__(self, test_client=None, quota_project_id=None):
self.test_client = test_client
+ self.quota_project_id = quota_project_id
def start_bundle(self):
- self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client)
+ self.bq_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client, quota_project_id=self.quota_project_id)
def process(self, table_reference):
_LOGGER.info("Deleting table %s", table_reference)
@@ -998,6 +1009,7 @@ def __init__(
self.test_client = test_client
self.schema = schema
+ self.quota_project_id = None
self._temp_file_format = temp_file_format or bigquery_tools.FileFormat.JSON
# If we have multiple destinations, then we will have multiple load jobs,
@@ -1207,7 +1219,8 @@ def _load_data(
additional_bq_parameters=self.additional_bq_parameters,
source_format=self._temp_file_format,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
load_job_name_pcv,
*self.schema_side_inputs).with_outputs(
TriggerLoadJobs.TEMP_TABLES,
@@ -1228,7 +1241,8 @@ def _load_data(
test_client=self.test_client,
additional_bq_parameters=self.additional_bq_parameters,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
schema_mod_job_name_pcv))
if self.write_disposition in ('WRITE_EMPTY', 'WRITE_TRUNCATE'):
@@ -1259,7 +1273,8 @@ def _load_data(
write_disposition=self.write_disposition,
test_client=self.test_client,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
copy_job_name_pcv,
pvalue.AsIter(schema_mod_job_ids_pc)).with_outputs(
TriggerCopyJobs.TRIGGER_DELETE_TEMP_TABLES, main='main'))
@@ -1275,7 +1290,9 @@ def _load_data(
| "RemoveTempTables/DeduplicateTables" >> beam.GroupByKey()
| "RemoveTempTables/GetTableNames" >> beam.Keys()
| "RemoveTempTables/Delete" >> beam.ParDo(
- DeleteTablesFn(self.test_client)))
+ DeleteTablesFn(
+ test_client=self.test_client,
+ quota_project_id=self.quota_project_id)))
# Load data directly to destination table
destination_load_job_ids_pc = (
@@ -1290,7 +1307,8 @@ def _load_data(
additional_bq_parameters=self.additional_bq_parameters,
source_format=self._temp_file_format,
step_name=step_name,
- load_job_project_id=self.load_job_project_id),
+ load_job_project_id=self.load_job_project_id,
+ quota_project_id=self.quota_project_id),
load_job_name_pcv,
*self.schema_side_inputs).with_outputs(
TriggerLoadJobs.ONGOING_JOBS, main='main')
@@ -1305,6 +1323,8 @@ def _load_data(
def expand(self, pcoll):
p = pcoll.pipeline
self.project = self.project or p.options.view_as(GoogleCloudOptions).project
+ self.quota_project_id = bigquery_tools._quota_project_id_from_options(
+ p.options)
try:
step_name = self.label
except AttributeError:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
index 7ae0ccc8e32f..7ccda7ecad8d 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
@@ -19,6 +19,7 @@
# pytype: skip-file
+import contextlib
import logging
import os
import secrets
@@ -1259,6 +1260,52 @@ def __call__(self):
assert_that(jobs, equal_to(expected_jobs), label='CheckJobs')
+class TestFileLoadsQuotaProject(_TestCaseWithTempDirCleanUp):
+ """Tests that the file loads DoFns attribute quota to the quota project."""
+ @parameterized.expand([
+ param(dofn=bqfl.UpdateDestinationSchema, quota_attr='_quota_project_id'),
+ param(dofn=bqfl.TriggerCopyJobs, quota_attr='quota_project_id'),
+ param(dofn=bqfl.TriggerLoadJobs, quota_attr='quota_project_id'),
+ param(dofn=bqfl.DeleteTablesFn, quota_attr='quota_project_id'),
+ ])
+ def test_start_bundle_passes_quota_project(self, dofn, quota_attr):
+ fn = dofn(quota_project_id='my-billing-project')
+ self.assertEqual(getattr(fn, quota_attr), 'my-billing-project')
+ with mock.patch.object(bigquery_tools, 'BigQueryWrapper') as mock_wrapper:
+ fn.start_bundle()
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project')
+
+ def test_expand_resolves_quota_project_from_options(self):
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ transform = bqfl.BigQueryBatchFileLoads(
+ 'project:dataset.table',
+ custom_gcs_temp_location=self._new_tempdir(),
+ test_client=mock.Mock(),
+ validate=False,
+ temp_file_format=bigquery_tools.FileFormat.JSON)
+ p = TestPipeline(options=options)
+ dofns = [
+ bqfl.UpdateDestinationSchema,
+ bqfl.TriggerCopyJobs,
+ bqfl.TriggerLoadJobs,
+ bqfl.DeleteTablesFn
+ ]
+ with contextlib.ExitStack() as stack:
+ mocks = [
+ stack.enter_context(
+ mock.patch.object(
+ dofn, '__init__', autospec=True, side_effect=dofn.__init__))
+ for dofn in dofns
+ ]
+ _ = p | beam.Create(_ELEMENTS) | transform
+ self.assertEqual(transform.quota_project_id, 'my-billing-project')
+ for dofn, m in zip(dofns, mocks):
+ self.assertTrue(m.call_args_list, dofn.__name__)
+ for c in m.call_args_list:
+ self.assertEqual(c.kwargs['quota_project_id'], 'my-billing-project')
+
+
class BigQueryFileLoadsIT(unittest.TestCase):
BIG_QUERY_SCHEMA = (
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
index 136b3cc56b7e..2fa95a3d1794 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
@@ -142,11 +142,13 @@ class _PassThroughThenCleanupTempDatasets(PTransform):
Utilizes readiness of PCollection to trigger DoFn.
"""
- def __init__(self, side_input=None):
+ def __init__(self, side_input=None, quota_project_id=None):
self.side_input = side_input
+ self.quota_project_id = quota_project_id
def expand(self, input):
pipeline_options = input.pipeline.options
+ quota_project_id = self.quota_project_id
class PassThrough(beam.DoFn):
def process(self, element):
@@ -154,8 +156,9 @@ def process(self, element):
class CleanUpProjects(beam.DoFn):
def process(self, unused_element, unused_signal, pipeline_details):
- bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(
- pipeline_options)
+ bq = bigquery_tools.BigQueryWrapper(
+ pipeline_options=pipeline_options,
+ quota_project_id=quota_project_id)
pipeline_details = pipeline_details[0]
if 'temp_table_ref' in pipeline_details.keys():
temp_table_ref = pipeline_details['temp_table_ref']
@@ -251,7 +254,7 @@ def _get_temp_dataset_project(self):
def start_bundle(self):
self.bq = bigquery_tools.BigQueryWrapper(
temp_dataset_id=self._get_temp_dataset_id(),
- client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options))
+ pipeline_options=self.options)
def process(self,
element: 'ReadFromBigQueryRequest') -> Iterable[BoundedSource]:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py
index 22b617fa8741..ae03104612c9 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py
@@ -53,6 +53,7 @@
from apache_beam.io.gcp.bigquery import _StreamToBigQuery
from apache_beam.io.gcp.bigquery_read_internal import _BigQueryReadSplit
from apache_beam.io.gcp.bigquery_read_internal import _JsonToDictCoder
+from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanupTempDatasets
from apache_beam.io.gcp.bigquery_read_internal import bigquery_export_destination_uri
from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR
from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
@@ -65,6 +66,7 @@
from apache_beam.io.gcp.tests.bigquery_matcher import BigQueryTableMatcher
from apache_beam.metrics.metric import Lineage
from apache_beam.options import value_provider
+from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.options.value_provider import RuntimeValueProvider
@@ -338,6 +340,10 @@ def test_repeatable_field_is_properly_converted(self):
class TestReadFromBigQuery(unittest.TestCase):
@classmethod
def setUpClass(cls):
+ cls.env_patch = mock.patch.dict(
+ os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'})
+ cls.env_patch.start()
+
class UserDefinedOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
@@ -351,6 +357,7 @@ def tearDown(self):
@classmethod
def tearDownClass(cls):
+ cls.env_patch.stop()
# Unset the option added in setupClass to avoid interfere with other tests.
# Force a gc so PipelineOptions.__subclass__() no longer contains it.
del cls.UserDefinedOptions
@@ -879,6 +886,343 @@ def test_expand_output_type_uses_query_schema(self):
mock_convert.assert_called_once_with(schema, None)
+@unittest.skipIf(
+ HttpError is None or gcp_bigquery is None,
+ 'GCP dependencies are not installed')
+class TestReadFromBigQueryQuotaProject(unittest.TestCase):
+ """Tests for quota_project_id in ReadFromBigQuery sources."""
+ def test_quota_project_id_from_pipeline_options(self):
+ """Test that quota_project_id is read from GoogleCloudOptions."""
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ gcp_options = options.view_as(GoogleCloudOptions)
+ self.assertEqual(gcp_options.quota_project_id, 'my-billing-project')
+
+ def test_quota_project_id_none_by_default_in_options(self):
+ """Test that quota_project_id is None by default in options."""
+ options = PipelineOptions([])
+ gcp_options = options.view_as(GoogleCloudOptions)
+ self.assertIsNone(gcp_options.quota_project_id)
+
+ def test_export_source_explicit_quota_project(self):
+ """Test that explicit quota_project_id is stored in
+ _CustomBigQuerySource."""
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ self.assertEqual(source.quota_project_id, 'my-billing-project')
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ def test_export_source_gets_quota_from_options(self):
+ """Test that _CustomBigQuerySource falls back to options for
+ quota_project_id."""
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ pipeline_options=options)
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ def test_export_source_explicit_overrides_options(self):
+ """Test that explicit quota_project_id overrides options."""
+ options = PipelineOptions(['--quota_project_id=options-project'])
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ pipeline_options=options,
+ quota_project_id='explicit-project')
+ self.assertEqual(source._get_quota_project_id(), 'explicit-project')
+
+ def test_storage_source_explicit_quota_project(self):
+ """Test that explicit quota_project_id is stored in
+ _CustomBigQueryStorageSource."""
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ self.assertEqual(source.quota_project_id, 'my-billing-project')
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ def test_storage_source_gets_quota_from_options(self):
+ """Test that _CustomBigQueryStorageSource falls back to options."""
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ pipeline_options=options)
+ self.assertEqual(source._get_quota_project_id(), 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_export_source_split_passes_quota_to_client(self, mock_bq_client):
+ """split() must pass the transform-level quota_project_id to the BigQuery
+ client, not only the --quota_project_id pipeline option."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ list(source.split(desired_bundle_size=0))
+ mock_bq_client.assert_called_once()
+ self.assertIsInstance(mock_bq_client.call_args.args[0], PipelineOptions)
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_storage_source_split_passes_quota_to_client(self, mock_bq_client):
+ """split() must pass the transform-level quota_project_id to the BigQuery
+ client, not only the --quota_project_id pipeline option."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ list(source.split(desired_bundle_size=0))
+ mock_bq_client.assert_called_once()
+ self.assertIsInstance(mock_bq_client.call_args.args[0], PipelineOptions)
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools, 'create_bq_storage_client')
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, 'get_table')
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_storage_source_split_passes_quota_to_storage_client(
+ self, unused_mock_bq_client, mock_get_table, mock_create_storage_client):
+ """split() must create the Storage Read client with the quota project
+ and pass it on to each stream source."""
+ mock_get_table.return_value = mock.Mock(numBytes=5)
+ stream = mock.Mock()
+ stream.name = 'projects/p/locations/l/sessions/s/streams/stream1'
+ mock_create_storage_client.return_value.create_read_session.return_value = (
+ mock.Mock(streams=[stream]))
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+
+ bundles = list(source.split(desired_bundle_size=1 << 20))
+
+ mock_create_storage_client.assert_called_once_with('my-billing-project')
+ self.assertEqual(len(bundles), 1)
+ self.assertEqual(bundles[0].source.quota_project_id, 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_export_source_estimate_size_passes_quota_to_client(
+ self, mock_bq_client):
+ """estimate_size() must attribute quota too, not just split()."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ source.estimate_size()
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client')
+ def test_storage_source_estimate_size_passes_quota_to_client(
+ self, mock_bq_client):
+ """estimate_size() must attribute quota too, not just split()."""
+ mock_bq_client.side_effect = RuntimeError('stop')
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ with self.assertRaises(RuntimeError):
+ source.estimate_size()
+ self.assertEqual(
+ mock_bq_client.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery.bigquery_schema_tools'
+ '.convert_to_usertype')
+ @mock.patch.object(bigquery_tools, 'BigQueryWrapper')
+ def test_output_type_schema_fetch_passes_quota_to_client(
+ self, mock_wrapper, mock_convert):
+ """The BEAM_ROW schema lookup is a BigQuery API call of this transform, so
+ it must be attributed to the quota project as well."""
+ mock_convert.return_value = beam.Map(lambda x: x)
+ transform = ReadFromBigQuery(
+ table='project:dataset.table',
+ output_type='BEAM_ROW',
+ quota_project_id='my-billing-project')
+ transform._expand_output_type(mock.MagicMock())
+
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project')
+
+ @mock.patch.object(bigquery_tools, 'BigQueryWrapper')
+ def test_direct_read_cleanup_passes_quota_to_client(self, mock_wrapper):
+ """The DIRECT_READ temp-dataset cleanup is a BigQuery API call of this
+ transform, so it must be attributed to the quota project as well."""
+ pipeline_details = {
+ 'project_id': 'project', 'bigquery_dataset_labels': {
+ 'k': 'v'
+ }
+ }
+ with TestPipeline() as p:
+ side_input = beam.pvalue.AsList(
+ p | 'Details' >> beam.Create([pipeline_details]))
+ _ = (
+ p | beam.Create([1])
+ | _PassThroughThenCleanupTempDatasets(
+ side_input, quota_project_id='my-billing-project'))
+
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project')
+ self.assertIsNotNone(mock_wrapper.call_args.kwargs['pipeline_options'])
+
+ @mock.patch('apache_beam.io.gcp.bigquery._PassThroughThenCleanupTempDatasets')
+ def test_direct_read_wires_quota_project_into_cleanup(self, mock_cleanup):
+ mock_cleanup.return_value = beam.Map(lambda x: x)
+ transform = ReadFromBigQuery(
+ table='project:dataset.table',
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ quota_project_id='my-billing-project')
+ p = TestPipeline()
+ transform._expand_direct_read(beam.pvalue.PBegin(p))
+
+ self.assertEqual(
+ mock_cleanup.call_args.kwargs['quota_project_id'], 'my-billing-project')
+
+ def test_quota_project_id_in_export_source_display_data(self):
+ """Test that quota_project_id appears in display data for export source."""
+ source = beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ items = {i.key: i.value for i in DisplayData.create_from(source).items}
+ self.assertEqual(items['quota_project_id'], 'my-billing-project')
+
+ def test_quota_project_id_in_direct_read_source_display_data(self):
+ """Test that quota_project_id appears in display data for DIRECT_READ."""
+ source = beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ quota_project_id='my-billing-project')
+ items = {i.key: i.value for i in DisplayData.create_from(source).items}
+ self.assertEqual(items['quota_project_id'], 'my-billing-project')
+
+ def test_quota_project_id_absent_from_display_data_when_not_set(self):
+ """Test that quota_project_id is dropped from display data when not set."""
+ options = PipelineOptions([])
+ sources = [
+ beam_bq._CustomBigQuerySource(
+ method=ReadFromBigQuery.Method.EXPORT,
+ table='project:dataset.table',
+ pipeline_options=options),
+ beam_bq._CustomBigQueryStorageSource(
+ method=ReadFromBigQuery.Method.DIRECT_READ,
+ table='project:dataset.table',
+ pipeline_options=options)
+ ]
+ for source in sources:
+ keys = {i.key for i in DisplayData.create_from(source).items}
+ self.assertNotIn('quota_project_id', keys)
+
+ def test_stream_source_stores_quota_project_id(self):
+ """Test that quota_project_id is stored in
+ _CustomBigQueryStorageStreamSource."""
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=True,
+ timeout=30.0,
+ quota_project_id='my-billing-project')
+ self.assertEqual(stream_source.quota_project_id, 'my-billing-project')
+
+ def test_stream_source_quota_project_id_none_by_default(self):
+ """Test that quota_project_id is None by default in stream source."""
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1')
+ self.assertIsNone(stream_source.quota_project_id)
+
+ def test_stream_source_split_preserves_quota_project_id(self):
+ """Test that split() preserves quota_project_id."""
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=True,
+ timeout=30.0,
+ quota_project_id='my-billing-project')
+ bundle = stream_source.split(desired_bundle_size=0)
+ self.assertEqual(bundle.source.quota_project_id, 'my-billing-project')
+ self.assertEqual(bundle.source.timeout, 30.0)
+ self.assertEqual(bundle.source.use_native_datetime, True)
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.create_bq_storage_client')
+ def test_stream_source_read_arrow_uses_quota_project(
+ self, mock_create_client):
+ """Test that read_arrow() uses create_bq_storage_client
+ with quota_project_id."""
+ mock_client = mock.MagicMock()
+ mock_create_client.return_value = mock_client
+ # Mock read_rows to return empty iterator
+ mock_client.read_rows.return_value.rows.return_value = iter([])
+
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=True,
+ quota_project_id='my-billing-project')
+ # Consume the iterator
+ list(stream_source.read_arrow())
+
+ mock_create_client.assert_called_once_with('my-billing-project')
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.create_bq_storage_client')
+ def test_stream_source_read_avro_uses_quota_project(self, mock_create_client):
+ """Test that read_avro() uses create_bq_storage_client
+ with quota_project_id."""
+ mock_client = mock.MagicMock()
+ mock_create_client.return_value = mock_client
+ # Mock read_rows to return empty iterator
+ mock_client.read_rows.return_value = iter([])
+
+ stream_source = beam_bq._CustomBigQueryStorageStreamSource(
+ read_stream_name='projects/p/locations/l/sessions/s/streams/stream1',
+ use_native_datetime=False,
+ quota_project_id='my-billing-project')
+ # Consume the iterator
+ list(stream_source.read_avro())
+
+ mock_create_client.assert_called_once_with('my-billing-project')
+
+
+@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
+class TestWriteToBigQueryQuotaProject(unittest.TestCase):
+ """Tests for quota_project_id in the BigQuery streaming insert path."""
+ @mock.patch.object(bigquery_tools, 'BigQueryWrapper')
+ def test_write_fn_start_bundle_passes_quota_project(self, mock_wrapper):
+ fn = beam_bq.BigQueryWriteFn(
+ batch_size=1, quota_project_id='my-billing-project')
+ fn.start_bundle()
+ self.assertEqual(
+ mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project')
+
+ def test_stream_to_bigquery_resolves_quota_project_from_options(self):
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ p = TestPipeline(options=options)
+ with mock.patch.object(
+ beam_bq.BigQueryWriteFn,
+ '__init__',
+ autospec=True,
+ side_effect=beam_bq.BigQueryWriteFn.__init__) as mock_write_fn:
+ _ = (
+ p
+ | beam.Create([])
+ | WriteToBigQuery(
+ 'project:dataset.table',
+ schema='a:STRING',
+ method=WriteToBigQuery.Method.STREAMING_INSERTS))
+ self.assertEqual(
+ mock_write_fn.call_args.kwargs['quota_project_id'],
+ 'my-billing-project')
+
+
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestBigQuerySink(unittest.TestCase):
def test_table_spec_display_data(self):
@@ -920,10 +1264,14 @@ def _cleanup_files(self):
os.remove('insert_calls2')
def setUp(self):
+ self.env_patch = mock.patch.dict(
+ os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'})
+ self.env_patch.start()
self._cleanup_files()
def tearDown(self):
self._cleanup_files()
+ self.env_patch.stop()
def test_noop_schema_parsing(self):
expected_table_schema = None
@@ -1481,6 +1829,13 @@ def test_copy_load_job_exception(self, exception_type, error_message):
HttpError is None or exceptions is None,
'GCP dependencies are not installed')
class BigQueryStreamingInsertsErrorHandling(unittest.TestCase):
+ def setUp(self):
+ self.env_patch = mock.patch.dict(
+ os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'})
+ self.env_patch.start()
+
+ def tearDown(self):
+ self.env_patch.stop()
# Running tests with a variety of exceptions from https://googleapis.dev
# /python/google-api-core/latest/_modules/google/api_core/exceptions.html.
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index cc1a16a659fe..d1412d0f855e 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -36,6 +36,7 @@
import re
import sys
import time
+import urllib.parse
import uuid
from json.decoder import JSONDecodeError
from typing import Optional
@@ -59,6 +60,7 @@
from apache_beam.metrics import monitoring_infos
from apache_beam.metrics.metric import Metrics
from apache_beam.options import value_provider
+from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.transforms import DoFn
from apache_beam.typehints.row_type import RowTypeConstraint
@@ -80,6 +82,11 @@
gcp_bigquery = None
pass
+try:
+ from google.cloud import bigquery_storage_v1 as bq_storage
+except ImportError:
+ bq_storage = None
+
try:
from orjson import dumps as fast_json_dumps
from orjson import loads as fast_json_loads
@@ -419,6 +426,58 @@ def _build_dataset_encryption_config(kms_key):
return bigquery.EncryptionConfiguration(kmsKeyName=kms_key)
+def _quota_project_id_from_options(pipeline_options):
+ """Returns --quota_project_id from pipeline options, if any."""
+ if pipeline_options is None:
+ return None
+ return pipeline_options.view_as(GoogleCloudOptions).quota_project_id
+
+
+def create_bq_storage_client(quota_project_id=None):
+ """Create a BigQueryReadClient with optional quota project.
+
+ Args:
+ quota_project_id: Optional GCP project ID to use for quota and billing.
+ The client applies it to the credentials it resolves itself.
+
+ Returns:
+ A BigQueryReadClient instance.
+ """
+ if not quota_project_id:
+ return bq_storage.BigQueryReadClient()
+ return bq_storage.BigQueryReadClient(
+ client_options={'quota_project_id': quota_project_id})
+
+
+class _HttpWithHeaders(object):
+ """Wraps an httplib2.Http to add fixed headers to requests to given hosts.
+
+ Sits below the credentials layer, so the headers are applied on every
+ attempt and take precedence over headers set by the credentials. The
+ credentials refresh their token through the same Http, so requests to other
+ hosts (token endpoint, metadata server, IAM Credentials API) are passed
+ through unchanged.
+ """
+ def __init__(self, http, headers, hosts):
+ object.__setattr__(self, '_http', http)
+ object.__setattr__(self, '_headers', headers)
+ object.__setattr__(self, '_hosts', frozenset(hosts))
+
+ def request(self, uri, method='GET', body=None, headers=None, **kwargs):
+ if urllib.parse.urlsplit(uri).hostname in self._hosts:
+ headers = dict(headers or {})
+ headers.update(self._headers)
+ return self._http.request(uri, method, body=body, headers=headers, **kwargs)
+
+ def __getattr__(self, name):
+ if name == '_http': # object created without __init__ (copy/unpickle)
+ raise AttributeError(name)
+ return getattr(self._http, name)
+
+ def __setattr__(self, name, value):
+ setattr(self._http, name, value)
+
+
class BigQueryWrapper(object):
"""BigQuery client wrapper with utilities for querying.
@@ -427,8 +486,10 @@ class BigQueryWrapper(object):
In addition, it offers various functions used both in sources and sinks
(e.g., find and create tables, query a table, etc.).
- Note that client parameter in constructor is only for testing purposes and
- should not be used in production code.
+ Note that the client and gcp_client parameters in the constructor are only
+ for testing purposes and should not be used in production code. Production
+ code should pass pipeline_options (and optionally quota_project_id) so that
+ every client is built here with the right quota attribution.
"""
# If updating following names, also update the corresponding pydocs in
@@ -438,16 +499,29 @@ class BigQueryWrapper(object):
HISTOGRAM_METRIC_LOGGER = MetricLogger()
- def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
- self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions())
- self.gcp_bq_client = client or gcp_bigquery.Client(
- client_info=ClientInfo(
- user_agent="apache-beam-%s" % apache_beam.__version__))
+ def __init__(
+ self,
+ client=None,
+ temp_dataset_id=None,
+ temp_table_ref=None,
+ quota_project_id=None,
+ *,
+ gcp_client=None,
+ pipeline_options=None):
+ self.quota_project_id = (
+ quota_project_id or _quota_project_id_from_options(pipeline_options))
+ self.client = client or BigQueryWrapper._bigquery_client(
+ pipeline_options or PipelineOptions(),
+ quota_project_id=self.quota_project_id)
+ # Test-only overrides. `client` also stands in for the google-cloud-bigquery
+ # client, which several tests rely on; pass `gcp_client` as well as `client`
+ # to override only that one.
+ self._gcp_bq_client = gcp_client or client
self._unique_row_id = 0
# For testing scenarios where we pass in a client we do not want a
# randomized prefix for row IDs.
- self._row_id_prefix = '' if client else uuid.uuid4()
+ self._row_id_prefix = '' if (client or gcp_client) else uuid.uuid4()
self._latency_histogram_metric = Metrics.histogram(
self.__class__,
'latency_histogram_ms',
@@ -474,6 +548,17 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
self.created_temp_dataset = False
+ @property
+ def gcp_bq_client(self):
+ """The google-cloud-bigquery client, created on first use.
+
+ Only the streaming insert path needs it, so it is not built eagerly.
+ """
+ if self._gcp_bq_client is None:
+ self._gcp_bq_client = BigQueryWrapper._gcp_bigquery_client(
+ quota_project_id=self.quota_project_id)
+ return self._gcp_bq_client
+
@property
def unique_row_id(self):
"""Returns a unique row ID (str) used to avoid multiple insertions.
@@ -1490,19 +1575,70 @@ def convert_row_to_dict(self, row, schema):
@staticmethod
def from_pipeline_options(pipeline_options: PipelineOptions):
- return BigQueryWrapper(
- client=BigQueryWrapper._bigquery_client(pipeline_options))
+ """Create a BigQueryWrapper from pipeline options.
+
+ Args:
+ pipeline_options: Pipeline options containing GCP configuration.
+ The quota_project_id is read from GoogleCloudOptions if set.
+ """
+ return BigQueryWrapper(pipeline_options=pipeline_options)
@staticmethod
- def _bigquery_client(pipeline_options: PipelineOptions):
+ def _bigquery_client(
+ pipeline_options: PipelineOptions, quota_project_id: str = None):
+ """Create a BigQuery API client from pipeline options.
+
+ Args:
+ pipeline_options: Pipeline options for credentials.
+ quota_project_id: Optional quota project ID. If not provided, will be
+ extracted from pipeline_options.
+ """
+ credentials = auth.get_service_credentials(pipeline_options)
+ # Use explicit quota_project_id if provided, otherwise get from options
+ quota_project_id = quota_project_id or _quota_project_id_from_options(
+ pipeline_options)
+ http = get_new_http()
+ if quota_project_id:
+ if credentials is None:
+ # Ignoring the request would silently bill a different project.
+ raise ValueError(
+ 'quota_project_id was set to %r, but no credentials were found to '
+ 'apply it to.' % quota_project_id)
+ # Send the quota project as a request header rather than deriving new
+ # credentials: credentials.with_quota_project() returns a copy without
+ # the cached token, and this keeps the shared credentials untouched.
+ # The header is set below the credentials layer so that it is applied on
+ # every attempt and wins over a quota project the credentials carry.
+ http = _HttpWithHeaders(
+ http, {'x-goog-user-project': quota_project_id},
+ hosts=[
+ urllib.parse.urlsplit(url).hostname for url in (
+ bigquery.BigqueryV2.BASE_URL,
+ bigquery.BigqueryV2.MTLS_BASE_URL)
+ ])
return bigquery.BigqueryV2(
- http=get_new_http(),
- credentials=auth.get_service_credentials(pipeline_options),
+ http=http,
+ credentials=credentials,
response_encoding='utf8',
additional_http_headers={
"user-agent": "apache-beam-%s" % apache_beam.__version__
})
+ @staticmethod
+ def _gcp_bigquery_client(quota_project_id: str = None):
+ """Create a google-cloud-bigquery Client with optional quota project.
+
+ The client applies the quota project to the credentials it resolves
+ itself, so no credentials are derived here.
+ """
+ client_options = None
+ if quota_project_id:
+ client_options = {'quota_project_id': quota_project_id}
+ return gcp_bigquery.Client(
+ client_options=client_options,
+ client_info=ClientInfo(
+ user_agent="apache-beam-%s" % apache_beam.__version__))
+
class RowAsDictJsonCoder(coders.Coder):
"""A coder for a table row (represented as a dict) to/from a JSON string.
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
index 9f8b78ee5c3b..de4f3e4066af 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -23,6 +23,7 @@
import json
import logging
import math
+import os
import re
import unittest
from typing import Optional
@@ -35,10 +36,12 @@
from parameterized import parameterized
import apache_beam as beam
+from apache_beam.io.gcp import bigquery_tools
from apache_beam.io.gcp import resource_identifiers
from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR
from apache_beam.io.gcp.bigquery_tools import AvroRowWriter
from apache_beam.io.gcp.bigquery_tools import BigQueryJobTypes
+from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
from apache_beam.io.gcp.bigquery_tools import JsonRowWriter
from apache_beam.io.gcp.bigquery_tools import RowAsDictJsonCoder
from apache_beam.io.gcp.bigquery_tools import beam_row_from_dict
@@ -299,14 +302,16 @@ def test_delete_dataset_retries_for_timeouts(self, patched_time_sleep):
@mock.patch('google.cloud._http.JSONConnection.http')
def test_user_agent_insert_all(
self, http_mock, patched_skip_get_credentials, patched_sleep):
- wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper()
- try:
- wrapper._insert_all_rows('p', 'd', 't', [{'name': 'any'}], None)
- except: # pylint: disable=bare-except
- # Ignore errors. The errors come from the fact that we did not mock
- # the response from the API, so the overall insert_all_rows call fails
- # soon after the BQ API is called.
- pass
+ # Set GOOGLE_CLOUD_PROJECT to ensure Client creation succeeds in test env
+ with mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'}):
+ wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper()
+ try:
+ wrapper._insert_all_rows('p', 'd', 't', [{'name': 'any'}], None)
+ except: # pylint: disable=bare-except
+ # Ignore errors. The errors come from the fact that we did not mock
+ # the response from the API, so the overall insert_all_rows call fails
+ # soon after the BQ API is called.
+ pass
call = http_mock.request.mock_calls[-2]
self.assertIn('apache-beam-', call[2]['headers']['User-Agent'])
@@ -1199,8 +1204,6 @@ def test_geography_in_bigquery_type_mapping(self):
def test_geography_field_conversion(self):
"""Test that GEOGRAPHY fields are converted correctly."""
- from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
-
# Create a mock field with GEOGRAPHY type
field = bigquery.TableFieldSchema()
field.type = 'GEOGRAPHY'
@@ -1322,8 +1325,6 @@ def test_geography_json_encoding(self):
def test_geography_with_special_characters(self):
"""Test GEOGRAPHY values with special characters and geometries."""
- from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
-
field = bigquery.TableFieldSchema()
field.type = 'GEOGRAPHY'
field.name = 'complex_geo'
@@ -1502,6 +1503,213 @@ def test_type_overrides_json_to_dict(self):
self.assertEqual(typehints_dict, [("data", Optional[dict])])
+@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
+class TestBigQueryWrapperQuotaProject(unittest.TestCase):
+ """Tests for quota_project_id in BigQueryWrapper."""
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client')
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client')
+ def test_quota_project_id_stored(self, mock_gcp_client, mock_bq_client):
+ """Test that quota_project_id is stored in BigQueryWrapper."""
+ mock_bq_client.return_value = mock.Mock()
+ mock_gcp_client.return_value = mock.Mock()
+
+ wrapper = BigQueryWrapper(quota_project_id='my-billing-project')
+ self.assertEqual(wrapper.quota_project_id, 'my-billing-project')
+
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client')
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client')
+ def test_from_pipeline_options_reads_quota_from_options(
+ self, mock_gcp_client, mock_bq_client):
+ """Test from_pipeline_options reads quota_project_id from
+ GoogleCloudOptions."""
+ from apache_beam.options.pipeline_options import PipelineOptions
+
+ mock_bq_client.return_value = mock.Mock()
+ mock_gcp_client.return_value = mock.Mock()
+
+ options = PipelineOptions(['--quota_project_id=my-billing-project'])
+ wrapper = BigQueryWrapper.from_pipeline_options(options)
+
+ self.assertEqual(wrapper.quota_project_id, 'my-billing-project')
+
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client')
+ @mock.patch(
+ 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client')
+ def test_from_pipeline_options_none_when_not_set(
+ self, mock_gcp_client, mock_bq_client):
+ """Test from_pipeline_options returns None when quota_project_id not set."""
+ from apache_beam.options.pipeline_options import PipelineOptions
+
+ mock_bq_client.return_value = mock.Mock()
+ mock_gcp_client.return_value = mock.Mock()
+
+ options = PipelineOptions([])
+ wrapper = BigQueryWrapper.from_pipeline_options(options)
+
+ self.assertIsNone(wrapper.quota_project_id)
+
+ def test_http_with_headers_sets_headers_and_delegates(self):
+ from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders
+ inner = mock.MagicMock()
+ inner.request.return_value = ('response', b'content')
+ inner.redirect_codes = {301, 302}
+ http = _HttpWithHeaders(
+ inner, {'x-goog-user-project': 'p'}, hosts=['example.com'])
+
+ result = http.request(
+ 'https://example.com',
+ method='GET',
+ headers={
+ 'x-goog-user-project': 'other', 'accept': 'json'
+ },
+ redirections=3)
+
+ self.assertEqual(result, ('response', b'content'))
+ inner.request.assert_called_once_with(
+ 'https://example.com',
+ 'GET',
+ body=None,
+ headers={
+ 'x-goog-user-project': 'p', 'accept': 'json'
+ },
+ redirections=3)
+ # Attribute reads and writes reach the wrapped http.
+ self.assertEqual(http.redirect_codes, {301, 302})
+ http.redirect_codes = {301}
+ self.assertEqual(inner.redirect_codes, {301})
+
+ def test_http_with_headers_survives_copy(self):
+ import copy
+
+ from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders
+ inner = mock.MagicMock()
+ http = copy.copy(
+ _HttpWithHeaders(
+ inner, {'x-goog-user-project': 'p'}, hosts=['example.com']))
+
+ http.request('https://example.com')
+ inner.request.assert_called_once_with(
+ 'https://example.com',
+ 'GET',
+ body=None,
+ headers={'x-goog-user-project': 'p'})
+
+ def test_http_with_headers_skips_other_hosts(self):
+ from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders
+ inner = mock.MagicMock()
+ http = _HttpWithHeaders(
+ inner, {'x-goog-user-project': 'p'}, hosts=['example.com'])
+
+ http.request('https://oauth2.googleapis.com/token', headers={'a': 'b'})
+ inner.request.assert_called_once_with(
+ 'https://oauth2.googleapis.com/token',
+ 'GET',
+ body=None,
+ headers={'a': 'b'})
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.bq_storage')
+ def test_create_bq_storage_client_with_quota_project(self, mock_bq_storage):
+ """The quota project is passed as a client option; the client applies it
+ to the credentials it resolves itself."""
+ bigquery_tools.create_bq_storage_client('my-billing-project')
+ mock_bq_storage.BigQueryReadClient.assert_called_once_with(
+ client_options={'quota_project_id': 'my-billing-project'})
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.bq_storage')
+ def test_create_bq_storage_client_without_quota_project(
+ self, mock_bq_storage):
+ """Without a quota project the client is created with defaults."""
+ bigquery_tools.create_bq_storage_client(None)
+ mock_bq_storage.BigQueryReadClient.assert_called_once_with()
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.auth.get_service_credentials')
+ def test_bigquery_client_raises_without_credentials(
+ self, mock_get_credentials):
+ """An explicit quota project must fail loudly rather than fall back to
+ apitools' own credential discovery, which would bill another project."""
+ from apache_beam.options.pipeline_options import PipelineOptions
+ mock_get_credentials.return_value = None
+
+ with self.assertRaisesRegex(ValueError, 'no credentials were found'):
+ BigQueryWrapper._bigquery_client(
+ PipelineOptions(), quota_project_id='requested-project')
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.get_new_http')
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.auth.get_service_credentials')
+ def test_bigquery_client_quota_project_wins_over_credentials(
+ self, mock_get_credentials, mock_get_new_http):
+ """The requested quota project must reach every BigQuery request,
+ including the retry after a 401 refresh, even when the credentials carry
+ their own, but not the token refresh requests."""
+ import httplib2
+ from google.auth import credentials as ga_credentials
+
+ from apache_beam.internal.gcp import auth
+ from apache_beam.options.pipeline_options import PipelineOptions
+
+ class CredentialsWithQuotaProject(ga_credentials.Credentials):
+ def __init__(self):
+ super().__init__()
+ self._quota_project_id = 'adc-project'
+
+ def refresh(self, request):
+ # Token refreshes go through the same Http as the API requests.
+ request('https://oauth2.googleapis.com/token', method='POST')
+ self.token = 'token'
+
+ seen_headers = []
+ token_request_headers = []
+ statuses = iter([401, 200])
+
+ class FakeHttp(object):
+ connections = {}
+ redirect_codes = set()
+
+ def request(self, uri, method='GET', body=None, headers=None, **kwargs):
+ if uri.startswith('https://oauth2.googleapis.com/'):
+ token_request_headers.append(dict(headers or {}))
+ return httplib2.Response({'status': 200}), b'{}'
+ seen_headers.append(dict(headers))
+ response = httplib2.Response({
+ 'status': next(statuses), 'content-type': 'application/json'
+ })
+ return response, b'{}'
+
+ mock_get_new_http.return_value = FakeHttp()
+ mock_get_credentials.return_value = auth._ApitoolsCredentialsAdapter(
+ CredentialsWithQuotaProject())
+
+ client = BigQueryWrapper._bigquery_client(
+ PipelineOptions(), quota_project_id='requested-project')
+ client.projects.List(bigquery.BigqueryProjectsListRequest())
+
+ self.assertEqual(len(seen_headers), 2)
+ for headers in seen_headers:
+ self.assertEqual(headers['x-goog-user-project'], 'requested-project')
+ self.assertEqual(headers['authorization'], 'Bearer token')
+ # The quota project must not be sent to the token endpoint.
+ self.assertTrue(token_request_headers)
+ for headers in token_request_headers:
+ self.assertNotIn('x-goog-user-project', headers)
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.gcp_bigquery.Client')
+ def test_gcp_bigquery_client_passes_quota_project(self, mock_client):
+ BigQueryWrapper._gcp_bigquery_client(quota_project_id='my-billing-project')
+ self.assertEqual(
+ mock_client.call_args.kwargs['client_options'],
+ {'quota_project_id': 'my-billing-project'})
+
+ @mock.patch('apache_beam.io.gcp.bigquery_tools.gcp_bigquery.Client')
+ def test_gcp_bigquery_client_no_quota_project_by_default(self, mock_client):
+ BigQueryWrapper._gcp_bigquery_client()
+ self.assertIsNone(mock_client.call_args.kwargs['client_options'])
+
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
unittest.main()
diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py
index 08ba34dabee4..f0406e2505bc 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -1197,6 +1197,17 @@ def _add_argparse_args(cls, parser):
action='store_true',
help='Throttling counter in GcsIO is enabled by default. Set '
'--no_gcsio_throttling_counter to avoid it.')
+ parser.add_argument(
+ '--quota_project_id',
+ default=None,
+ help='GCP project ID used for quota and billing attribution of Google '
+ 'Cloud API requests (sets the x-goog-user-project header), if '
+ 'different from the project the data resides in. Currently honored '
+ 'by BigQuery IO transforms (ReadFromBigQuery, ReadAllFromBigQuery, '
+ 'WriteToBigQuery with FILE_LOADS or STREAMING_INSERTS, '
+ 'ReadBigQueryChangeHistory); STORAGE_WRITE_API (Java-backed) and '
+ 'other GCP IOs do not use it yet. The credentials used must have the '
+ 'serviceusage.services.use permission on this project.')
parser.add_argument(
'--enable_gcsio_blob_generation',
default=False,
From 707b951e5e04015d195a62b5ab51d2088391fc7f Mon Sep 17 00:00:00 2001
From: GAURAV THAKUR
Date: Mon, 19 Oct 2026 08:44:19 +0530
Subject: [PATCH 025/108] Preserve HTTP errors when quota response is not JSON
(#40631)
---
sdks/python/apache_beam/utils/retry.py | 2 +-
sdks/python/apache_beam/utils/retry_test.py | 18 ++++++++++++++++++
2 files changed, 19 insertions(+), 1 deletion(-)
diff --git a/sdks/python/apache_beam/utils/retry.py b/sdks/python/apache_beam/utils/retry.py
index 03e42829b843..063446472fe5 100644
--- a/sdks/python/apache_beam/utils/retry.py
+++ b/sdks/python/apache_beam/utils/retry.py
@@ -183,7 +183,7 @@ def retry_on_server_errors_timeout_or_quota_issues_filter(exception):
if not isinstance(content, dict):
content = json.loads(exception.content)
return content["error"]["errors"][0]["reason"] in _RETRYABLE_REASONS
- except (KeyError, IndexError, TypeError) as e:
+ except (KeyError, IndexError, TypeError, ValueError) as e:
_LOGGER.warning(
"Could not determine if HttpError is transient. "
"Will not retry: %s",
diff --git a/sdks/python/apache_beam/utils/retry_test.py b/sdks/python/apache_beam/utils/retry_test.py
index 05cc6e6a7f40..8d9e006e00c5 100644
--- a/sdks/python/apache_beam/utils/retry_test.py
+++ b/sdks/python/apache_beam/utils/retry_test.py
@@ -141,6 +141,24 @@ def test_with_http_error_that_should_be_retried(self):
500)
self.assertEqual(len(self.clock.calls), 10)
+ @unittest.skipIf(HttpError is None, 'google-apitools is not installed')
+ def test_quota_filter_preserves_http_error_with_non_json_content(self):
+ error = HttpError({'status': '403'},
+ 'upstream error',
+ 'https://example.test')
+
+ @retry.with_exponential_backoff(
+ clock=self.clock,
+ retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter
+ )
+ def fail():
+ raise error
+
+ with self.assertRaises(HttpError) as raised:
+ fail()
+ self.assertIs(raised.exception, error)
+ self.assertEqual(self.clock.calls, [])
+
def test_with_explicit_initial_delay(self):
self.assertRaises(
NotImplementedError,
From 5fc82a67af69229a4a1597d91de8acd5b98b3569 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rados=C5=82aw=20Stankiewicz?=
Date: Mon, 19 Oct 2026 09:57:13 +0200
Subject: [PATCH 026/108] [Spanner Change Streams] Match single-RPC deadline to
total timeout in SpannerIO.ReadChangeStream (#40616)
Set initialRpcTimeout and maxRpcTimeout to 5 minutes to match totalTimeout in SpannerIO.ReadChangeStream.buildChangeStreamSpannerConfig() instead of using a 1-minute per-RPC deadline.
---
.../beam/sdk/io/gcp/spanner/SpannerIO.java | 4 +-
.../SpannerIOReadChangeStreamTest.java | 52 +++++++++++++++++++
2 files changed, 54 insertions(+), 2 deletions(-)
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
index cc573686357a..c14bb6dfdcdc 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
@@ -2538,8 +2538,8 @@ SpannerConfig buildChangeStreamSpannerConfig() {
.setExecuteStreamingSqlRetrySettings(
RetrySettings.newBuilder()
.setTotalTimeout(org.threeten.bp.Duration.ofMinutes(5))
- .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(1))
- .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(1))
+ .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(5))
+ .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(5))
.build())
.build();
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java
index 13fd343f5c83..3c003e5cbd4b 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadChangeStreamTest.java
@@ -26,6 +26,8 @@
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.StatusCode.Code;
import com.google.auth.Credentials;
import com.google.cloud.Timestamp;
import com.google.cloud.spanner.DatabaseClient;
@@ -41,6 +43,7 @@
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.io.gcp.spanner.changestreams.MetadataSpannerConfigFactory;
import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -175,6 +178,55 @@ public void testSetDirectedReadOptions() {
.getReplicaSelections(0)
.getLocation());
}
+
+ @Test
+ public void testDefaultRetrySettings() {
+ assertNull(spannerConfig.getRetryableCodes());
+ assertNull(spannerConfig.getExecuteStreamingSqlRetrySettings());
+ SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig();
+ SpannerConfig metadataSpannerConfig =
+ MetadataSpannerConfigFactory.create(
+ changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE);
+ ImmutableSet expectedRetryableCodes = ImmutableSet.of(Code.UNAVAILABLE, Code.ABORTED);
+ RetrySettings expectedRetrySettings =
+ RetrySettings.newBuilder()
+ .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(5))
+ .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(5))
+ .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(5))
+ .build();
+ assertEquals(expectedRetryableCodes, changeStreamSpannerConfig.getRetryableCodes());
+ assertEquals(expectedRetryableCodes, metadataSpannerConfig.getRetryableCodes());
+ assertEquals(
+ expectedRetrySettings, changeStreamSpannerConfig.getExecuteStreamingSqlRetrySettings());
+ assertEquals(
+ expectedRetrySettings, metadataSpannerConfig.getExecuteStreamingSqlRetrySettings());
+ }
+
+ @Test
+ public void testCustomRetrySettings() {
+ ImmutableSet customRetryableCodes = ImmutableSet.of(Code.UNAVAILABLE);
+ RetrySettings customRetrySettings =
+ RetrySettings.newBuilder()
+ .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(10))
+ .setInitialRpcTimeout(org.threeten.bp.Duration.ofMinutes(10))
+ .setMaxRpcTimeout(org.threeten.bp.Duration.ofMinutes(10))
+ .build();
+ spannerConfig =
+ spannerConfig
+ .withRetryableCodes(customRetryableCodes)
+ .withExecuteStreamingSqlRetrySettings(customRetrySettings);
+ readChangeStream = readChangeStream.withSpannerConfig(spannerConfig);
+ SpannerConfig changeStreamSpannerConfig = readChangeStream.buildChangeStreamSpannerConfig();
+ SpannerConfig metadataSpannerConfig =
+ MetadataSpannerConfigFactory.create(
+ changeStreamSpannerConfig, TEST_METADATA_INSTANCE, TEST_METADATA_DATABASE);
+ assertEquals(customRetryableCodes, changeStreamSpannerConfig.getRetryableCodes());
+ assertEquals(customRetryableCodes, metadataSpannerConfig.getRetryableCodes());
+ assertEquals(
+ customRetrySettings, changeStreamSpannerConfig.getExecuteStreamingSqlRetrySettings());
+ assertEquals(
+ customRetrySettings, metadataSpannerConfig.getExecuteStreamingSqlRetrySettings());
+ }
}
/** Parameterized tests for Dialect and Partition Mode combinations. */
From 10d3b8d316d76bc19b8f19508063388dd3896728 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rados=C5=82aw=20Stankiewicz?=
Date: Mon, 19 Oct 2026 10:00:47 +0200
Subject: [PATCH 027/108] Fix temp table cleanup for continuous read in
BigQueryIO (#37535)
* Fix temp table cleanup for continuous read in BigQueryIO
* Fix missing coder for CleanupOperationMessage side outputs
* Fix BQ streaming input tasks: move Javadoc, fix empty table edge case, and add test case.
---
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 242 ++++++++++++++++--
.../bigquery/BigQueryStorageStreamSource.java | 21 +-
.../bigquery/CleanupTempTableDoFnTest.java | 153 +++++++++++
3 files changed, 396 insertions(+), 20 deletions(-)
create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
index 0cefea4639fd..07fa1a320e92 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -23,6 +23,7 @@
import static org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.RECORDING_ROUTER;
import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects.firstNonNull;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
@@ -117,9 +118,14 @@
import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
import org.apache.beam.sdk.schemas.ProjectionProducer;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
+import org.apache.beam.sdk.transforms.DoFn.StateId;
+import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -141,6 +147,7 @@
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
+import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.Row;
@@ -957,11 +964,18 @@ abstract Builder setBadRecordErrorHandler(
DynamicRead() {}
class CreateBoundedSourceForTable
- extends DoFn, BigQueryStorageStreamSource> {
+ extends DoFn<
+ KV, KV>> {
+
+ private final TupleTag> cleanupInfoTag;
+
+ CreateBoundedSourceForTable(TupleTag> cleanupInfoTag) {
+ this.cleanupInfoTag = cleanupInfoTag;
+ }
@ProcessElement
public void processElement(
- OutputReceiver> receiver,
+ MultiOutputReceiver receiver,
@Element KV kv,
PipelineOptions options)
throws Exception {
@@ -986,7 +1000,14 @@ public void processElement(
// shards
long desiredChunkSize = getDesiredChunkSize(options, output);
List> split = output.split(desiredChunkSize, options);
- split.stream().forEach(source -> receiver.output(source));
+ split.stream()
+ .forEach(
+ source ->
+ receiver
+ .get(
+ new TupleTag>>(
+ "mainOutput"))
+ .output(KV.of(kv.getKey(), source)));
} else {
// run query
BigQueryStorageQuerySource querySource =
@@ -1023,7 +1044,21 @@ public void processElement(
// shards
long desiredChunkSize = getDesiredChunkSize(options, output);
List> split = output.split(desiredChunkSize, options);
- split.stream().forEach(source -> receiver.output(source));
+ split.stream()
+ .forEach(
+ source ->
+ receiver
+ .get(
+ new TupleTag>>(
+ "mainOutput"))
+ .output(KV.of(kv.getKey(), source.withFromQuery())));
+ boolean datasetCreatedByBeam = getQueryTempDataset() == null;
+ CleanupInfo cleanupInfo =
+ new CleanupInfo(
+ queryResultTable.getTableReference(), datasetCreatedByBeam, split.size());
+ receiver
+ .get(cleanupInfoTag)
+ .output(KV.of(kv.getKey(), CleanupOperationMessage.initialize(cleanupInfo)));
}
}
@@ -1037,6 +1072,9 @@ private long getDesiredChunkSize(
@Override
public PCollection expand(PCollection input) {
TupleTag rowTag = new TupleTag<>();
+ TupleTag>> streamTag = new TupleTag<>("mainOutput");
+ TupleTag> cleanupInfoTag = new TupleTag<>();
+
PCollection> addJobId =
input
.apply(
@@ -1051,21 +1089,193 @@ public String apply(BigQueryDynamicReadDescriptor input) {
.apply("Checkpoint", Redistribute.byKey());
PCollectionTuple resultTuple =
- addJobId
- .apply("Create streams", ParDo.of(new CreateBoundedSourceForTable()))
+ addJobId.apply(
+ "Create streams",
+ ParDo.of(new CreateBoundedSourceForTable(cleanupInfoTag))
+ .withOutputTags(streamTag, TupleTagList.of(cleanupInfoTag)));
+
+ PCollection>> streams =
+ resultTuple
+ .get(streamTag)
.setCoder(
- SerializableCoder.of(new TypeDescriptor>() {}))
- .apply("Redistribute", Redistribute.arbitrarily())
- .apply(
- "Read Streams with storage read api",
- ParDo.of(
- new TypedRead.ReadTableSource(
- rowTag, getParseFn(), getBadRecordRouter()))
- .withOutputTags(rowTag, TupleTagList.of(BAD_RECORD_TAG)));
+ KvCoder.of(
+ StringUtf8Coder.of(),
+ SerializableCoder.of(
+ new TypeDescriptor>() {})))
+ .apply("Redistribute", Redistribute.arbitrarily());
+
+ PCollectionTuple readResultTuple =
+ streams.apply(
+ "Read Streams with storage read api",
+ ParDo.of(
+ new ReadDynamicStreamSource(
+ rowTag, getParseFn(), getBadRecordRouter(), cleanupInfoTag))
+ .withOutputTags(rowTag, TupleTagList.of(BAD_RECORD_TAG).and(cleanupInfoTag)));
+
+ PCollection> cleanupMessages1 =
+ resultTuple
+ .get(cleanupInfoTag)
+ .setCoder(
+ KvCoder.of(
+ StringUtf8Coder.of(), SerializableCoder.of(CleanupOperationMessage.class)));
+
+ PCollection> cleanupMessages2 =
+ readResultTuple
+ .get(cleanupInfoTag)
+ .setCoder(
+ KvCoder.of(
+ StringUtf8Coder.of(), SerializableCoder.of(CleanupOperationMessage.class)));
+
+ PCollectionList.of(cleanupMessages1)
+ .and(cleanupMessages2)
+ .apply(Flatten.pCollections())
+ .apply("CleanupTempTables", ParDo.of(new CleanupTempTableDoFn(getBigQueryServices())));
+
getBadRecordErrorHandler()
.addErrorCollection(
- resultTuple.get(BAD_RECORD_TAG).setCoder(BadRecord.getCoder(input.getPipeline())));
- return resultTuple.get(rowTag).setCoder(getOutputCoder());
+ readResultTuple
+ .get(BAD_RECORD_TAG)
+ .setCoder(BadRecord.getCoder(input.getPipeline())));
+ return readResultTuple.get(rowTag).setCoder(getOutputCoder());
+ }
+ }
+
+ static class CleanupInfo implements Serializable {
+ private final String projectId;
+ private final String datasetId;
+ private final String tableId;
+ private final boolean datasetCreatedByBeam;
+ private final int totalStreams;
+
+ public CleanupInfo(TableReference tableRef, boolean datasetCreatedByBeam, int totalStreams) {
+ this.projectId = tableRef.getProjectId();
+ this.datasetId = tableRef.getDatasetId();
+ this.tableId = tableRef.getTableId();
+ this.datasetCreatedByBeam = datasetCreatedByBeam;
+ this.totalStreams = totalStreams;
+ }
+
+ public TableReference getTableReference() {
+ return new TableReference()
+ .setProjectId(projectId)
+ .setDatasetId(datasetId)
+ .setTableId(tableId);
+ }
+
+ public boolean isDatasetCreatedByBeam() {
+ return datasetCreatedByBeam;
+ }
+
+ public int getTotalStreams() {
+ return totalStreams;
+ }
+ }
+
+ static class CleanupOperationMessage implements Serializable {
+ private final @Nullable CleanupInfo cleanupInfo;
+ private final boolean isStreamCompletion;
+
+ private CleanupOperationMessage(@Nullable CleanupInfo cleanupInfo, boolean isStreamCompletion) {
+ this.cleanupInfo = cleanupInfo;
+ this.isStreamCompletion = isStreamCompletion;
+ }
+
+ public static CleanupOperationMessage streamComplete() {
+ return new CleanupOperationMessage(null, true);
+ }
+
+ public static CleanupOperationMessage initialize(CleanupInfo cleanupInfo) {
+ return new CleanupOperationMessage(cleanupInfo, false);
+ }
+
+ public @Nullable CleanupInfo getCleanupInfo() {
+ return cleanupInfo;
+ }
+
+ public boolean isStreamCompletion() {
+ return isStreamCompletion;
+ }
+ }
+
+ static class CleanupTempTableDoFn extends DoFn, Void> {
+ private final BigQueryServices bqServices;
+ private static final Logger LOG = LoggerFactory.getLogger(CleanupTempTableDoFn.class);
+
+ @StateId("cleanupInfo")
+ private final StateSpec> cleanupInfoSpec = StateSpecs.value();
+
+ @StateId("completedStreams")
+ private final StateSpec> completedStreamsSpec = StateSpecs.value();
+
+ CleanupTempTableDoFn(BigQueryServices bqServices) {
+ this.bqServices = bqServices;
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV element,
+ @StateId("cleanupInfo") ValueState cleanupInfoState,
+ @StateId("completedStreams") ValueState completedStreamsState,
+ PipelineOptions options)
+ throws Exception {
+
+ CleanupOperationMessage msg = element.getValue();
+ CleanupInfo cleanupInfo = cleanupInfoState.read();
+ int completed = firstNonNull(completedStreamsState.read(), 0);
+
+ if (msg.isStreamCompletion()) {
+ completed += 1;
+ completedStreamsState.write(completed);
+ } else {
+ cleanupInfoState.write(Preconditions.checkNotNull(msg.getCleanupInfo()));
+ cleanupInfo = msg.getCleanupInfo();
+ }
+
+ if (cleanupInfo != null && completed == cleanupInfo.getTotalStreams()) {
+ TableReference tempTable = cleanupInfo.getTableReference();
+ try (DatasetService datasetService =
+ bqServices.getDatasetService(options.as(BigQueryOptions.class))) {
+ LOG.info("Deleting temporary table with query results {}", tempTable);
+ datasetService.deleteTable(tempTable);
+ if (cleanupInfo.isDatasetCreatedByBeam()) {
+ LOG.info("Deleting temporary dataset with query results {}", tempTable.getDatasetId());
+ datasetService.deleteDataset(tempTable.getProjectId(), tempTable.getDatasetId());
+ }
+ } catch (Exception e) {
+ LOG.warn("Failed to delete temporary BigQuery table {}", tempTable, e);
+ }
+ cleanupInfoState.clear();
+ completedStreamsState.clear();
+ }
+ }
+ }
+
+ private static class ReadDynamicStreamSource
+ extends DoFn>, T> {
+ private final TypedRead.ReadTableSource readTableSource;
+ private final TupleTag> streamCompletionTag;
+
+ ReadDynamicStreamSource(
+ TupleTag rowTag,
+ SerializableFunction parseFn,
+ BadRecordRouter badRecordRouter,
+ TupleTag> streamCompletionTag) {
+ this.readTableSource = new TypedRead.ReadTableSource<>(rowTag, parseFn, badRecordRouter);
+ this.streamCompletionTag = streamCompletionTag;
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV> element,
+ MultiOutputReceiver receiver,
+ PipelineOptions options)
+ throws Exception {
+ readTableSource.processElement(element.getValue(), receiver, options);
+ if (element.getValue().getFromQuery()) {
+ receiver
+ .get(streamCompletionTag)
+ .output(KV.of(element.getKey(), CleanupOperationMessage.streamComplete()));
+ }
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
index 1bc0b2895d9a..6376af7d816f 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
@@ -77,7 +77,13 @@ public static BigQueryStorageStreamSource create(
toJsonString(Preconditions.checkArgumentNotNull(tableSchema, "tableSchema")),
parseFn,
outputCoder,
- bqServices);
+ bqServices,
+ false);
+ }
+
+ public BigQueryStorageStreamSource withFromQuery() {
+ return new BigQueryStorageStreamSource<>(
+ readSession, readStream, jsonTableSchema, parseFn, outputCoder, bqServices, true);
}
@Override
@@ -106,13 +112,13 @@ public int hashCode() {
*/
public BigQueryStorageStreamSource fromExisting(ReadStream newReadStream) {
return new BigQueryStorageStreamSource<>(
- readSession, newReadStream, jsonTableSchema, parseFn, outputCoder, bqServices);
+ readSession, newReadStream, jsonTableSchema, parseFn, outputCoder, bqServices, fromQuery);
}
public BigQueryStorageStreamSource fromExisting(
SerializableFunction parseFn) {
return new BigQueryStorageStreamSource<>(
- readSession, readStream, jsonTableSchema, parseFn, outputCoder, bqServices);
+ readSession, readStream, jsonTableSchema, parseFn, outputCoder, bqServices, fromQuery);
}
private final ReadSession readSession;
@@ -121,6 +127,7 @@ public BigQueryStorageStreamSource fromExisting(
private final SerializableFunction parseFn;
private final Coder outputCoder;
private final BigQueryServices bqServices;
+ private final boolean fromQuery;
private BigQueryStorageStreamSource(
ReadSession readSession,
@@ -128,13 +135,19 @@ private BigQueryStorageStreamSource(
String jsonTableSchema,
SerializableFunction parseFn,
Coder outputCoder,
- BigQueryServices bqServices) {
+ BigQueryServices bqServices,
+ boolean fromQuery) {
this.readSession = Preconditions.checkArgumentNotNull(readSession, "readSession");
this.readStream = Preconditions.checkArgumentNotNull(readStream, "stream");
this.jsonTableSchema = Preconditions.checkArgumentNotNull(jsonTableSchema, "jsonTableSchema");
this.parseFn = Preconditions.checkArgumentNotNull(parseFn, "parseFn");
this.outputCoder = Preconditions.checkArgumentNotNull(outputCoder, "outputCoder");
this.bqServices = Preconditions.checkArgumentNotNull(bqServices, "bqServices");
+ this.fromQuery = fromQuery;
+ }
+
+ public boolean getFromQuery() {
+ return fromQuery;
}
@Override
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java
new file mode 100644
index 000000000000..b1fa3dc8e576
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import com.google.api.services.bigquery.model.TableReference;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.CleanupInfo;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.CleanupOperationMessage;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.CleanupTempTableDoFn;
+import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices;
+import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class CleanupTempTableDoFnTest {
+
+ @Rule public transient TestPipeline p = TestPipeline.create();
+
+ @Before
+ public void setUp() {
+ FakeDatasetService.setUp();
+ }
+
+ @Test
+ public void testCleanupTempTableDoFn() throws Exception {
+ FakeDatasetService fakeDatasetService = new FakeDatasetService();
+ FakeBigQueryServices fakeBqServices =
+ new FakeBigQueryServices().withDatasetService(fakeDatasetService);
+
+ TableReference tableRef =
+ new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+
+ fakeDatasetService.createDataset(
+ tableRef.getProjectId(), tableRef.getDatasetId(), "", "", null);
+ fakeDatasetService.createTable(
+ new com.google.api.services.bigquery.model.Table().setTableReference(tableRef));
+
+ assertTrue(
+ fakeDatasetService.getDataset(tableRef.getProjectId(), tableRef.getDatasetId()) != null);
+ assertTrue(fakeDatasetService.getTable(tableRef) != null);
+
+ CleanupInfo cleanupInfo = new CleanupInfo(tableRef, true, 2);
+
+ PCollection> input =
+ p.apply(
+ Create.of(
+ KV.of("job1", CleanupOperationMessage.initialize(cleanupInfo)),
+ KV.of("job1", CleanupOperationMessage.streamComplete()),
+ KV.of("job1", CleanupOperationMessage.streamComplete())));
+
+ input.apply(ParDo.of(new CleanupTempTableDoFn(fakeBqServices)));
+
+ p.run().waitUntilFinish();
+
+ // The dataset is deleted, so getDataset and getTable should throw a 404 Exception or return
+ // null
+ try {
+ fakeDatasetService.getDataset(tableRef.getProjectId(), tableRef.getDatasetId());
+ fail("Dataset should have been deleted");
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+
+ try {
+ TableReference deletedRef =
+ new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+ Object table = fakeDatasetService.getTable(deletedRef);
+ assertTrue(table == null);
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+ }
+
+ @Test
+ public void testCleanupTempTableDoFnWithZeroStreams() throws Exception {
+ FakeDatasetService fakeDatasetService = new FakeDatasetService();
+ FakeBigQueryServices fakeBqServices =
+ new FakeBigQueryServices().withDatasetService(fakeDatasetService);
+
+ TableReference tableRef =
+ new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+
+ fakeDatasetService.createDataset(
+ tableRef.getProjectId(), tableRef.getDatasetId(), "", "", null);
+ fakeDatasetService.createTable(
+ new com.google.api.services.bigquery.model.Table().setTableReference(tableRef));
+
+ assertTrue(
+ fakeDatasetService.getDataset(tableRef.getProjectId(), tableRef.getDatasetId()) != null);
+ assertTrue(fakeDatasetService.getTable(tableRef) != null);
+
+ CleanupInfo cleanupInfo = new CleanupInfo(tableRef, true, 0);
+
+ PCollection> input =
+ p.apply(
+ "CreateInputWithZeroStreams",
+ Create.of(KV.of("job1", CleanupOperationMessage.initialize(cleanupInfo))));
+
+ input.apply("CleanupWithZeroStreams", ParDo.of(new CleanupTempTableDoFn(fakeBqServices)));
+
+ p.run().waitUntilFinish();
+
+ try {
+ fakeDatasetService.getDataset(tableRef.getProjectId(), tableRef.getDatasetId());
+ fail("Dataset should have been deleted");
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+
+ try {
+ TableReference deletedRef =
+ new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
+ Object table = fakeDatasetService.getTable(deletedRef);
+ assertTrue(table == null);
+ } catch (Exception e) {
+ assertTrue(
+ e.getMessage().contains("Tried to get a dataset")
+ || e.getMessage().contains("Not Found"));
+ }
+ }
+}
From 49d41adc89549aea325378e0007873b7b352c2b7 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 19 Oct 2026 07:47:37 -0400
Subject: [PATCH 028/108] Bump google.golang.org/api from 0.302.0 to 0.304.0 in
/sdks (#40650)
Bumps [google.golang.org/api](https://github.com/googleapis/google-api-go-client) from 0.302.0 to 0.304.0.
- [Release notes](https://github.com/googleapis/google-api-go-client/releases)
- [Changelog](https://github.com/googleapis/google-api-go-client/blob/main/CHANGES.md)
- [Commits](https://github.com/googleapis/google-api-go-client/compare/v0.302.0...v0.304.0)
---
updated-dependencies:
- dependency-name: google.golang.org/api
dependency-version: 0.304.0
dependency-type: direct:production
update-type: version-update:semver-minor
...
Signed-off-by: dependabot[bot]
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
sdks/go.mod | 2 +-
sdks/go.sum | 4 ++--
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/sdks/go.mod b/sdks/go.mod
index d60a9db1651f..f6b6d047f8f1 100644
--- a/sdks/go.mod
+++ b/sdks/go.mod
@@ -60,7 +60,7 @@ require (
golang.org/x/sync v0.24.0
golang.org/x/sys v0.49.0
golang.org/x/text v0.43.0
- google.golang.org/api v0.302.0
+ google.golang.org/api v0.304.0
google.golang.org/genproto v0.0.0-20260715232425-e75dac1f907d // indirect
google.golang.org/grpc v1.84.0
google.golang.org/protobuf v1.36.12
diff --git a/sdks/go.sum b/sdks/go.sum
index 1fc97c94e394..52ab74e2243e 100644
--- a/sdks/go.sum
+++ b/sdks/go.sum
@@ -1315,8 +1315,8 @@ google.golang.org/api v0.69.0/go.mod h1:boanBiw+h5c3s+tBPgEzLDRHfFLWV0qXxRHz3ws7
google.golang.org/api v0.70.0/go.mod h1:Bs4ZM2HGifEvXwd50TtW70ovgJffJYw2oRCOFU/SkfA=
google.golang.org/api v0.71.0/go.mod h1:4PyU6e6JogV1f9eA4voyrTY2batOLdgZ5qZ5HOCc4j8=
google.golang.org/api v0.74.0/go.mod h1:ZpfMZOVRMywNyvJFeqL9HRWBgAuRfSjJFpe9QtRRyDs=
-google.golang.org/api v0.302.0 h1:/eENsEv3+GE0QiEmXjdaLI84Mt76LNjGz4vp4E7sRVw=
-google.golang.org/api v0.302.0/go.mod h1:dcjRWAy1QdjGb3Z67Qu0gdBWH2gL8Qoy/RI2ZVuvdeE=
+google.golang.org/api v0.304.0 h1:2R1dYSiIVQRGcKN901u9yPFEdJu5h90oGY5X4TcgbmA=
+google.golang.org/api v0.304.0/go.mod h1:dcjRWAy1QdjGb3Z67Qu0gdBWH2gL8Qoy/RI2ZVuvdeE=
google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
google.golang.org/appengine v1.5.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
From e4bc725b42d388cf19f6781f6ee196f57cb52289 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 19 Oct 2026 07:48:20 -0400
Subject: [PATCH 029/108] Bump zizmorcore/zizmor-action from 0.6.4 to 0.7.1
(#40652)
Bumps [zizmorcore/zizmor-action](https://github.com/zizmorcore/zizmor-action) from 0.6.4 to 0.7.1.
- [Release notes](https://github.com/zizmorcore/zizmor-action/releases)
- [Commits](https://github.com/zizmorcore/zizmor-action/compare/cc914d7f3750a2d13d75c7f184a1060aa0e9d482...85b2b9fb8e3325e7a23d7d0a9d7ba5ae0f999e3e)
---
updated-dependencies:
- dependency-name: zizmorcore/zizmor-action
dependency-version: 0.7.1
dependency-type: direct:production
update-type: version-update:semver-minor
...
Signed-off-by: dependabot[bot]
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
.github/workflows/beam_PreCommit_GHA.yml | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/.github/workflows/beam_PreCommit_GHA.yml b/.github/workflows/beam_PreCommit_GHA.yml
index 04dd4d098783..0fabd369611f 100644
--- a/.github/workflows/beam_PreCommit_GHA.yml
+++ b/.github/workflows/beam_PreCommit_GHA.yml
@@ -96,7 +96,7 @@ jobs:
- name: Validate GHA Allowlist
uses: apache/infrastructure-actions/allowlist-check@main # zizmor: ignore[unpinned-uses]
- name: Run zizmor
- uses: zizmorcore/zizmor-action@cc914d7f3750a2d13d75c7f184a1060aa0e9d482 # v0.6.4
+ uses: zizmorcore/zizmor-action@85b2b9fb8e3325e7a23d7d0a9d7ba5ae0f999e3e # v0.7.1
with:
advanced-security: true
- name: run GHA PreCommit script
From 5fa8c2c878e553ee12aee8ca114a5124a9187ab8 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 19 Oct 2026 08:06:57 -0400
Subject: [PATCH 030/108] Bump github/codeql-action from 4.38.3 to 4.39.0
(#40651)
Bumps [github/codeql-action](https://github.com/github/codeql-action) from 4.38.3 to 4.39.0.
- [Release notes](https://github.com/github/codeql-action/releases)
- [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md)
- [Commits](https://github.com/github/codeql-action/compare/v4.38.3...v4.39.0)
---
updated-dependencies:
- dependency-name: github/codeql-action
dependency-version: 4.39.0
dependency-type: direct:production
update-type: version-update:semver-minor
...
Signed-off-by: dependabot[bot]
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
.github/workflows/codeql.yml | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml
index 883f2ca317a5..804f38d09922 100644
--- a/.github/workflows/codeql.yml
+++ b/.github/workflows/codeql.yml
@@ -154,7 +154,7 @@ jobs:
# Initializes the CodeQL tools for scanning.
- name: Initialize CodeQL
- uses: github/codeql-action/init@v4.38.3
+ uses: github/codeql-action/init@v4.39.0
with:
languages: ${{ matrix.language }}
build-mode: ${{ matrix.build-mode }}
@@ -193,6 +193,6 @@ jobs:
fi
- name: Perform CodeQL Analysis
- uses: github/codeql-action/analyze@v4.38.3
+ uses: github/codeql-action/analyze@v4.39.0
with:
category: "/language:${{matrix.language}}"
From 8648b89152356fa843dea19a0ad8a20e1bf7522b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rados=C5=82aw=20Stankiewicz?=
Date: Mon, 19 Oct 2026 15:44:08 +0200
Subject: [PATCH 031/108] Update users.yml (#40654)
---
infra/iam/users.yml | 7 ++++++-
1 file changed, 6 insertions(+), 1 deletion(-)
diff --git a/infra/iam/users.yml b/infra/iam/users.yml
index 6c37b4c1fef0..29e55c624cb6 100644
--- a/infra/iam/users.yml
+++ b/infra/iam/users.yml
@@ -745,6 +745,11 @@
- role: roles/cloudfunctions.invoker
- role: roles/datastore.user
- role: roles/storage.objectViewer
+- username: radoslaws
+ email: radoslaws@google.com
+ member_type: user
+ permissions:
+ - role: roles/editor
- username: rajkumargupta
email: rajkumargupta@google.com
member_type: user
@@ -1236,4 +1241,4 @@
member_type: user
permissions:
- role: projects/apache-beam-testing/roles/beam_viewer
- - role: projects/apache-beam-testing/roles/beam_writer
\ No newline at end of file
+ - role: projects/apache-beam-testing/roles/beam_writer
From 590abac114546ec33cedff26e44f37857150eb71 Mon Sep 17 00:00:00 2001
From: raams95
Date: Mon, 19 Oct 2026 09:53:25 -0500
Subject: [PATCH 032/108] [Python] Avoid ZeroDivisionError in
ConcatRangeTracker for zero weights (#40549)
* [Python] Avoid ZeroDivisionError in ConcatRangeTracker for zero weights
When every SourceBundle passed to a ConcatSource had weight 0, the
average used to fill in missing weights was also 0, so computing the
cumulative weights divided by zero. Treat zero weights like unknown
weights when computing that average.
Fixes #21316
* Add CHANGES.md entry
* Test zero-weight ConcatSource through a pipeline read
---------
Co-authored-by: JanakiRamChittineni <341781521+raams95@users.noreply.github.com>
Co-authored-by: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com>
---
CHANGES.md | 1 +
sdks/python/apache_beam/io/concat_source.py | 4 ++--
sdks/python/apache_beam/io/concat_source_test.py | 9 +++++++++
3 files changed, 12 insertions(+), 2 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index c10b7262209c..1ec1dbe80864 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -115,6 +115,7 @@
* (Python) `Duration` built from float seconds now rounds to the nearest microsecond instead of truncating, which could lose a microsecond ([#40263](https://github.com/apache/beam/issues/40263)).
* (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that directly follows an escaped delimiter, which merged two records into one ([#40459](https://github.com/apache/beam/issues/40459)).
* (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)).
+* (Python) `ConcatSource` no longer raises `ZeroDivisionError` when every source bundle reports a weight of 0, e.g. bundles smaller than 1 MB ([#21316](https://github.com/apache/beam/issues/21316)).
* (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)).
* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
diff --git a/sdks/python/apache_beam/io/concat_source.py b/sdks/python/apache_beam/io/concat_source.py
index 35ae6fe27317..bd6b7396d5d3 100644
--- a/sdks/python/apache_beam/io/concat_source.py
+++ b/sdks/python/apache_beam/io/concat_source.py
@@ -135,8 +135,8 @@ def _compute_cumulative_weights(source_bundles):
min_diff = 1e-5
# For the computation below, we need weights for all sources.
# Substitute average weights for those whose weights are
- # unspecified (or 1.0 for everything if none are known).
- known = [s.weight for s in source_bundles if s.weight is not None]
+ # unspecified or zero (or 1.0 for everything if none are known).
+ known = [s.weight for s in source_bundles if s.weight]
avg = sum(known) / len(known) if known else 1.0
weights = [s.weight or avg for s in source_bundles]
diff --git a/sdks/python/apache_beam/io/concat_source_test.py b/sdks/python/apache_beam/io/concat_source_test.py
index efa24b3975fc..3cd169df5069 100644
--- a/sdks/python/apache_beam/io/concat_source_test.py
+++ b/sdks/python/apache_beam/io/concat_source_test.py
@@ -171,6 +171,15 @@ def test_position_at_fration(self):
self.assertEqual(range_tracker.position_at_fraction(.601), (2, 17))
self.assertEqual(range_tracker.position_at_fraction(1), (3, None))
+ def test_read_zero_weight_bundles_in_pipeline(self):
+ source = ConcatSource([
+ iobase.SourceBundle(0, RangeSource(0, 4), None, None),
+ iobase.SourceBundle(0, RangeSource(4, 8), None, None),
+ ])
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | beam.io.Read(source)
+ assert_that(pcoll, equal_to(list(range(8))))
+
def test_empty_source(self):
read_all = source_test_utils.read_from_source
From 14575dcd143edd5f6f22fdff703671bda4aaf457 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rados=C5=82aw=20Stankiewicz?=
Date: Mon, 19 Oct 2026 16:55:00 +0200
Subject: [PATCH 033/108] Update api key for tour of beam (#40655)
---
learning/tour-of-beam/frontend/lib/firebase_options.dart | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/learning/tour-of-beam/frontend/lib/firebase_options.dart b/learning/tour-of-beam/frontend/lib/firebase_options.dart
index e64b594af433..fb6c35ce66f8 100644
--- a/learning/tour-of-beam/frontend/lib/firebase_options.dart
+++ b/learning/tour-of-beam/frontend/lib/firebase_options.dart
@@ -53,7 +53,7 @@ class DefaultFirebaseOptions {
}
static const FirebaseOptions web = FirebaseOptions(
- apiKey: 'AIzaSyAg7ZLslQRrEhwVVVzZb1OFMRMHL8kNL38',
+ apiKey: 'AIzaSyA4MXDLiEEdY-Bm_Iyxrg3W_cmdCSa4lvg',
authDomain: 'apache-beam-testing.firebaseapp.com',
projectId: 'apache-beam-testing',
storageBucket: 'apache-beam-testing.appspot.com',
From 2c60b1996e7c671104f98314ec70101f382c4158 Mon Sep 17 00:00:00 2001
From: Vitaly Terentyev
Date: Mon, 19 Oct 2026 18:55:39 +0400
Subject: [PATCH 034/108] Fix postgres database for grafana (#40656)
---
.../grafana/provisioning/datasources/beampostgresql.yaml | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/.test-infra/metrics/grafana/provisioning/datasources/beampostgresql.yaml b/.test-infra/metrics/grafana/provisioning/datasources/beampostgresql.yaml
index dd9a5d764a6e..496914bbd415 100644
--- a/.test-infra/metrics/grafana/provisioning/datasources/beampostgresql.yaml
+++ b/.test-infra/metrics/grafana/provisioning/datasources/beampostgresql.yaml
@@ -29,8 +29,8 @@ datasources:
secureJsonData:
password: ${PSQL_DB_DBPWD}
user: ${PSQL_DB_DBUSERNAME}
- database: ${PSQL_DB_DBNAME}
jsonData:
+ database: ${PSQL_DB_DBNAME}
sslmode: disable
editable: false
From bd102581f7f8fd9f446a22ff490cf272be9f4077 Mon Sep 17 00:00:00 2001
From: Abdelrahman Ibrahim
Date: Mon, 19 Oct 2026 18:29:50 +0300
Subject: [PATCH 035/108] Run portable ValidatesRunner tests one at a time
(#40632)
* run portable ValidatesRunner tests one at a time
* delay flink teardown
---
.../beam_PostCommit_Java_PVR_Flink_Batch.json | 2 +-
...m_PostCommit_Java_PVR_Flink_Streaming.json | 2 +-
.../FlinkExecutableStageContextFactory.java | 51 ++++++++++++++-----
.../flink/job-server/flink_job_server.gradle | 2 +
.../FlinkExecutableStageContextFactory.java | 51 ++++++++++++++-----
5 files changed, 78 insertions(+), 30 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json
index 6f486a1f8c14..73e12ca8961b 100644
--- a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json
+++ b/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json
@@ -1,6 +1,6 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"https://github.com/apache/beam/pull/32440": "test new datastream runner for batch",
- "modification": 2,
+ "modification": 8,
"https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
}
diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json
index 626e4fea5c8b..ebd004055319 100644
--- a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json
+++ b/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json
@@ -1,6 +1,6 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
- "modification": 2,
+ "modification": 3,
"https://github.com/apache/beam/pull/32440": "test new datastream runner for batch",
"https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
}
diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java
index 3f42eb93e4e6..afcfb9d67300 100644
--- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java
+++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java
@@ -22,7 +22,9 @@
import org.apache.beam.runners.fnexecution.control.DefaultExecutableStageContext;
import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
import org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory;
+import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
/** Singleton class that contains one {@link ExecutableStageContext.Factory} per job. */
@@ -30,10 +32,6 @@ public class FlinkExecutableStageContextFactory implements ExecutableStageContex
private static final FlinkExecutableStageContextFactory instance =
new FlinkExecutableStageContextFactory();
- // This map should only ever have a single element, as each job will have its own
- // classloader and therefore its own instance of FlinkExecutableStageContextFactory. This
- // code supports multiple JobInfos in order to provide a sensible implementation of
- // Factory.get(JobInfo), which in theory could be called with different JobInfos.
private static final ConcurrentMap jobFactories =
new ConcurrentHashMap<>();
@@ -48,16 +46,41 @@ public ExecutableStageContext get(JobInfo jobInfo) {
ExecutableStageContext.Factory jobFactory =
jobFactories.computeIfAbsent(
jobInfo.jobId(),
- k -> {
- return ReferenceCountingExecutableStageContextFactory.create(
- DefaultExecutableStageContext::create,
- // Clean up context immediately if its class is not loaded on Flink parent
- // classloader.
- (caller) ->
- caller.getClass().getClassLoader()
- != StreamExecutionEnvironment.class.getClassLoader());
- });
-
+ k ->
+ ReferenceCountingExecutableStageContextFactory.create(
+ info -> new JobScopedExecutableStageContext(info, k),
+ // Release immediately when the factory is not on the Flink parent classloader
+ // (so cleanup still works after that classloader is discarded). Otherwise keep
+ // the default environment cache delay so DoFn @Teardown can finish first.
+ caller ->
+ caller.getClass().getClassLoader()
+ != StreamExecutionEnvironment.class.getClassLoader()));
return jobFactory.get(jobInfo);
}
+
+ private static class JobScopedExecutableStageContext implements ExecutableStageContext {
+ private final DefaultExecutableStageContext delegate;
+ private final String jobId;
+
+ private JobScopedExecutableStageContext(JobInfo jobInfo, String jobId) {
+ this.delegate = DefaultExecutableStageContext.create(jobInfo);
+ this.jobId = jobId;
+ }
+
+ @Override
+ public StageBundleFactory getStageBundleFactory(ExecutableStage executableStage) {
+ return delegate.getStageBundleFactory(executableStage);
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ delegate.close();
+ } finally {
+ // Drop the per-job factory entry so a long-lived JVM that runs many jobs does not
+ // accumulate one entry per finished job.
+ jobFactories.remove(jobId);
+ }
+ }
+ }
}
diff --git a/runners/flink/job-server/flink_job_server.gradle b/runners/flink/job-server/flink_job_server.gradle
index 6f39b67abf74..edf55807503f 100644
--- a/runners/flink/job-server/flink_job_server.gradle
+++ b/runners/flink/job-server/flink_job_server.gradle
@@ -129,6 +129,7 @@ runShadow {
jvmArgs += ["-Dorg.slf4j.simpleLogger.defaultLogLevel=${project.property('logLevel')}"]
}
+
def portableValidatesRunnerTask(String name, String mode, boolean checkpointing, boolean docker) {
def pipelineOptions = [
// Limit resource consumption via parallelism
@@ -257,6 +258,7 @@ def flinkJobServerJvmArgs() {
[validatesPortableRunnerDocker, validatesPortableRunnerBatchDataSet, validatesPortableRunnerBatch, validatesPortableRunnerStreaming, validatesPortableRunnerStreamingCheckpoint].each { task ->
task.configure {
jvmArgs += flinkJobServerJvmArgs()
+ systemProperty "beam.validatesRunner.parallelThreads", "1"
}
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java
index a2e3bb3142bf..5a11336add14 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java
@@ -22,7 +22,9 @@
import org.apache.beam.runners.fnexecution.control.DefaultExecutableStageContext;
import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
import org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory;
+import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
import org.apache.flink.api.java.ExecutionEnvironment;
/** Singleton class that contains one {@link ExecutableStageContext.Factory} per job. */
@@ -30,10 +32,6 @@ public class FlinkExecutableStageContextFactory implements ExecutableStageContex
private static final FlinkExecutableStageContextFactory instance =
new FlinkExecutableStageContextFactory();
- // This map should only ever have a single element, as each job will have its own
- // classloader and therefore its own instance of FlinkExecutableStageContextFactory. This
- // code supports multiple JobInfos in order to provide a sensible implementation of
- // Factory.get(JobInfo), which in theory could be called with different JobInfos.
private static final ConcurrentMap jobFactories =
new ConcurrentHashMap<>();
@@ -48,16 +46,41 @@ public ExecutableStageContext get(JobInfo jobInfo) {
ExecutableStageContext.Factory jobFactory =
jobFactories.computeIfAbsent(
jobInfo.jobId(),
- k -> {
- return ReferenceCountingExecutableStageContextFactory.create(
- DefaultExecutableStageContext::create,
- // Clean up context immediately if its class is not loaded on Flink parent
- // classloader.
- (caller) ->
- caller.getClass().getClassLoader()
- != ExecutionEnvironment.class.getClassLoader());
- });
-
+ k ->
+ ReferenceCountingExecutableStageContextFactory.create(
+ info -> new JobScopedExecutableStageContext(info, k),
+ // Release immediately when the factory is not on the Flink parent classloader
+ // (so cleanup still works after that classloader is discarded). Otherwise keep
+ // the default environment cache delay so DoFn @Teardown can finish first.
+ caller ->
+ caller.getClass().getClassLoader()
+ != ExecutionEnvironment.class.getClassLoader()));
return jobFactory.get(jobInfo);
}
+
+ private static class JobScopedExecutableStageContext implements ExecutableStageContext {
+ private final DefaultExecutableStageContext delegate;
+ private final String jobId;
+
+ private JobScopedExecutableStageContext(JobInfo jobInfo, String jobId) {
+ this.delegate = DefaultExecutableStageContext.create(jobInfo);
+ this.jobId = jobId;
+ }
+
+ @Override
+ public StageBundleFactory getStageBundleFactory(ExecutableStage executableStage) {
+ return delegate.getStageBundleFactory(executableStage);
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ delegate.close();
+ } finally {
+ // Drop the per-job factory entry so a long-lived JVM that runs many jobs does not
+ // accumulate one entry per finished job.
+ jobFactories.remove(jobId);
+ }
+ }
+ }
}
From 62357960fc3114da572ef549287da4566699f770 Mon Sep 17 00:00:00 2001
From: tvalentyn
Date: Mon, 19 Oct 2026 08:35:00 -0700
Subject: [PATCH 036/108] Remove a race with another file write that is not
essential to test presence of duplicates. (#40596)
* Remove a race with another file write that is not essential to test presence of duplicates.
* Replace removed functionality.
---
sdks/python/apache_beam/io/fileio_test.py | 10 +---------
1 file changed, 1 insertion(+), 9 deletions(-)
diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py
index 1c650653804c..3bb250a34bd2 100644
--- a/sdks/python/apache_beam/io/fileio_test.py
+++ b/sdks/python/apache_beam/io/fileio_test.py
@@ -366,13 +366,6 @@ def test_without_deduplication(self):
file = self._create_temp_file(dir=tempdir)
# Add file twice, since it will be matched for every interval
files += [file, file]
- # Add file name that will be created mid-pipeline
- files.append(FileSystems.join(tempdir, 'extra'))
-
- def _create_extra_file(element):
- writer = FileSystems.create(FileSystems.join(tempdir, 'extra'))
- writer.close()
- return element.path
with TestPipeline() as p:
match_continiously = (
@@ -383,8 +376,7 @@ def _create_extra_file(element):
has_deduplication=False,
start_timestamp=start,
stop_timestamp=stop)
- | beam.Map(_create_extra_file))
-
+ | beam.Map(lambda element: element.path))
assert_that(match_continiously, equal_to(files))
def test_match_updated_files(self):
From a8ee2b84d37d014c3de8934b2e3ef8702260a2b2 Mon Sep 17 00:00:00 2001
From: Tobias Kaymak
Date: Mon, 19 Oct 2026 17:37:24 +0200
Subject: [PATCH 037/108] [Spark 4] Group the streaming stateful ParDo by
encoded key bytes to support null keys (#40516)
* [Spark 4] Group the streaming stateful ParDo by encoded key bytes to support null keys
A null Beam key became a null Spark grouping key and transformWithState failed with INTERNAL_ERROR_TWS. Grouping by the coder encoded key bytes is null safe and matches Beam key equality. Part of #36841.
* Retrigger CI, flaky StatefulStreamingParDoEvaluatorTest on Spark 3.4
* Retrigger CI, flaky legacy Spark streaming tests in Spark Versions
* Retrigger CI, KafkaToPubsubE2ETest container startup flake in PreCommit Java
* Use the structured streaming CoderHelpers in the stateful ParDo
* Merge the structured streaming CoderHelpers into the common one
* Decode with the common CoderHelpers in BeamMicroBatchSourceTest
---
...ostCommit_Java_ValidatesRunner_Spark4.json | 2 +-
...idatesRunner_SparkStructuredStreaming.json | 2 +-
.../io/streaming/BeamReaderCache.java | 2 +-
.../io/streaming/UnboundedSourceDataset.java | 2 +-
.../batch/GroupByKeyTranslatorBatch.java | 2 +-
.../translation/helpers/EncoderHelpers.java | 1 +
.../streaming/ReadUnboundedTranslator.java | 2 +-
.../StatefulParDoStreamingTranslator.java | 10 ++--
.../state/BeamStatefulProcessor.java | 13 ++--
.../streaming/BeamMicroBatchSourceTest.java | 12 ++--
.../streaming/StatefulParDoStreamingTest.java | 36 +++++++++++
.../batch/GroupByKeyTranslatorBatch.java | 2 +-
.../batch/functions/SideInputValues.java | 2 +-
.../translation/helpers/CoderHelpers.java | 59 -------------------
.../translation/helpers/EncoderHelpers.java | 1 +
15 files changed, 64 insertions(+), 84 deletions(-)
delete mode 100644 runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/CoderHelpers.java
diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json
index c537844dc84a..e0266d62f2e0 100644
--- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json
+++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
- "modification": 3
+ "modification": 4
}
diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json
index bd93fd4e9881..c0d5a77072b2 100644
--- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json
+++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json
@@ -12,5 +12,5 @@
"https://github.com/apache/beam/pull/40103": "noting that PR #40103 should run this test",
"https://github.com/apache/beam/issues/40427": "Spark ValidatesRunner tests run serially",
"https://github.com/apache/beam/pull/40281": "noting that PR #40281 should run this test",
- "https://github.com/apache/beam/issues/36841": "Spark 4 streaming watermark lag"
+ "https://github.com/apache/beam/issues/36841": "Spark 4 streaming null keys"
}
diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java
index 1eb31e766e78..e961ff2f3807 100644
--- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java
+++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java
@@ -27,7 +27,7 @@
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.LongSupplier;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java
index 0368f95ba6db..33106824c107 100644
--- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java
+++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java
@@ -29,9 +29,9 @@
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamReaderCache.CachedReader;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java
index 7caf06cb38fd..2b930c579d00 100644
--- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java
+++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java
@@ -18,13 +18,13 @@
package org.apache.beam.runners.spark.structuredstreaming.translation.batch;
import static org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY;
+import static org.apache.beam.runners.spark.coders.CoderHelpers.toByteArray;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGlobalGroupBy;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGroupByWindow;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.explodeWindowedKey;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueKey;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueValue;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.windowedKV;
-import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers.toByteArray;
import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.collectionEncoder;
import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.encoderOf;
import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.kvEncoder;
diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
index 173b4653a19b..4cd92da0a15f 100644
--- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
+++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
@@ -36,6 +36,7 @@
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java
index fb1691108b71..eb1421b45540 100644
--- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java
+++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java
@@ -19,10 +19,10 @@
import static org.apache.spark.sql.functions.col;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset;
import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
index 1709ec570038..199570ae64a1 100644
--- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
+++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.runners.spark.structuredstreaming.translation.streaming;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
import org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory;
@@ -35,7 +36,7 @@
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.Dataset;
-import org.apache.spark.sql.Encoder;
+import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.streaming.OutputMode;
import org.apache.spark.sql.streaming.TimeMode;
@@ -57,7 +58,7 @@ protected void translate(ParDo.MultiOutput, OutputT> transform, Context
DoFn, OutputT> doFn = transform.getFn();
KvCoder inputCoder = (KvCoder) input.getCoder();
- Encoder keyEnc = cxt.keyEncoderOf(inputCoder);
+ Coder keyCoder = inputCoder.getKeyCoder();
TupleTag mainOutputTag = transform.getMainOutputTag();
PCollection output = cxt.getOutput(mainOutputTag);
Coder outputCoder = output.getCoder();
@@ -71,10 +72,11 @@ protected void translate(ParDo.MultiOutput, OutputT> transform, Context
new BeamStatefulProcessor<>(
doFn, runnerFactory, inputCoder, windowing, cxt.getOptionsSupplier(), metrics);
- MapFunction>, K> keyFn = v -> v.getValue().getKey();
+ MapFunction>, byte[]> keyFn =
+ v -> CoderHelpers.toByteArray(v.getValue().getKey(), keyCoder);
Dataset> result =
cxt.getDataset(input)
- .groupByKey(keyFn, keyEnc)
+ .groupByKey(keyFn, Encoders.BINARY())
.transformWithState(
processor,
TimeMode.EventTime(),
diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java
index e359ed6bd58b..2649f883af3f 100644
--- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java
+++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java
@@ -41,6 +41,7 @@
import org.apache.beam.runners.spark.structuredstreaming.translation.batch.StatefulTaskRunner;
import org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.ListCoder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.TimeDomain;
@@ -68,7 +69,7 @@
/** Runs a stateful {@link DoFn} within Spark 4 Structured Streaming {@code transformWithState}. */
public class BeamStatefulProcessor
- extends StatefulProcessor>, WindowedValue> {
+ extends StatefulProcessor>, WindowedValue> {
private static final String BEAM_STATE = "beamState";
private static final String BEAM_TIMERS = "beamTimers";
@@ -76,7 +77,7 @@ public class BeamStatefulProcessor
private final DoFn, OutputT> doFn;
private final DoFnRunnerFactory, OutputT> runnerFactory;
- private final Coder> inputCoder;
+ private final KvCoder inputCoder;
private final WindowingStrategy windowingStrategy;
private final Supplier optionsSupplier;
private final MetricsAccumulator metrics;
@@ -93,7 +94,7 @@ public class BeamStatefulProcessor
public BeamStatefulProcessor(
DoFn, OutputT> doFn,
DoFnRunnerFactory, OutputT> runnerFactory,
- Coder> inputCoder,
+ KvCoder inputCoder,
WindowingStrategy windowingStrategy,
Supplier optionsSupplier,
MetricsAccumulator metrics) {
@@ -138,7 +139,8 @@ private SparkTimerInternals timerInternals(TimerValues timerValues) {
@Override
public Iterator> handleInputRows(
- K key, Iterator>> rows, TimerValues timerValues) {
+ byte[] keyBytes, Iterator>> rows, TimerValues timerValues) {
+ K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder());
List> outputs = new ArrayList<>();
SparkTimerInternals timerInternals = timerInternals(timerValues);
@@ -159,7 +161,8 @@ public Iterator> handleInputRows(
@Override
public Iterator> handleExpiredTimer(
- K key, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) {
+ byte[] keyBytes, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) {
+ K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder());
List> outputs = new ArrayList<>();
SparkTimerInternals timerInternals = timerInternals(timerValues);
timerInternals.advanceWatermark();
diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java
index 3b2e36bf5ef7..559deb808036 100644
--- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java
+++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java
@@ -46,6 +46,7 @@
import javax.annotation.Nullable;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.spark.StreamingTest;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule;
import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamInputPartition;
@@ -60,7 +61,6 @@
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.hadoop.conf.Configuration;
@@ -209,7 +209,7 @@ public void testReadsElementsFromUnboundedSource() throws Exception {
List values = new ArrayList<>();
for (List batch : batches(tag)) {
for (Row row : batch) {
- WindowedValue value = CoderUtils.decodeFromByteArray(CODER, row.getAs(COL_PAYLOAD));
+ WindowedValue value = CoderHelpers.fromByteArray(row.getAs(COL_PAYLOAD), CODER);
values.add(value.getValue());
assertEquals(
value.getTimestamp().getMillis(), row.getAs(COL_EVENT_TS).getTime());
@@ -593,11 +593,7 @@ private static List values(List> batches) {
}
private static String decode(byte[] payload) {
- try {
- return CoderUtils.decodeFromByteArray(CODER, payload).getValue();
- } catch (IOException e) {
- throw new IllegalStateException(e);
- }
+ return CoderHelpers.fromByteArray(payload, CODER).getValue();
}
private static List nonEmptySizes(List> batches) {
@@ -748,7 +744,7 @@ private static int committedPosition(File checkpointDir, int shard) throws IOExc
new BeamSourceCheckpoint(sourceDir(checkpointDir).getAbsolutePath(), new Configuration());
byte[] coded = checkpoint.readMark(shard, epoch);
assertNotNull("no mark at committed epoch " + epoch + " for shard " + shard, coded);
- return CoderUtils.decodeFromByteArray(TestUnboundedSource.MARK_CODER, coded).next;
+ return CoderHelpers.fromByteArray(coded, TestUnboundedSource.MARK_CODER).next;
}
// ---------------------------------------------------------------------------------------------
diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java
index 2ff0bfc2bf05..2d3c127d72bc 100644
--- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java
+++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java
@@ -33,6 +33,7 @@
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
@@ -43,6 +44,7 @@
import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
@@ -69,12 +71,14 @@ public class StatefulParDoStreamingTest implements Serializable {
private static final String TAG_BATCHES = "stateful-batches";
private static final String TAG_RESTART = "stateful-restart";
private static final String TAG_CHAINED = "stateful-chained";
+ private static final String TAG_NULL_KEY = "stateful-null-key";
@After
public void tearDown() {
TestUnboundedSource.forget(TAG_BATCHES);
TestUnboundedSource.forget(TAG_RESTART);
TestUnboundedSource.forget(TAG_CHAINED);
+ TestUnboundedSource.forget(TAG_NULL_KEY);
}
private static Pipeline pipeline(
@@ -167,6 +171,25 @@ public void testChainedStatefulParDoTimerOutputSurvives() throws Exception {
assertTrue(collected.contains("timer-output"));
}
+ @Test
+ public void testNullKey() throws Exception {
+ String collectorId = StreamingTestUtils.newCollectorId(TAG_NULL_KEY);
+ SparkStructuredStreamingPipelineOptions options =
+ StreamingTestUtils.streamingOptions(tempFolder);
+ options.setMaxRecordsPerBatch(1L);
+ Pipeline pipeline = Pipeline.create(options);
+ pipeline
+ .apply("Read", Read.from(new TestUnboundedSource(TAG_NULL_KEY, 1, 5)))
+ .apply("WithKeys", WithKeys.of((Void) null))
+ .setCoder(KvCoder.of(VoidCoder.of(), StringUtf8Coder.of()))
+ .apply("StatefulParDo", ParDo.of(new NullKeyCountFn()))
+ .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+ assertEquals(PipelineResult.State.DONE, StreamingTestUtils.run(pipeline).getState());
+ Set collected = StreamingTestUtils.collected(collectorId);
+ assertTrue(collected.contains(5));
+ }
+
private static final class KeySelectorFn extends DoFn> {
@ProcessElement
public void process(@Element String element, OutputReceiver> out) {
@@ -301,4 +324,17 @@ public void process(
out.output(element.getValue());
}
}
+
+ private static final class NullKeyCountFn extends DoFn, Integer> {
+ @StateId("count")
+ private final StateSpec> countSpec = StateSpecs.value();
+
+ @ProcessElement
+ public void process(
+ @StateId("count") ValueState countState, OutputReceiver out) {
+ int count = (countState.read() == null ? 0 : countState.read()) + 1;
+ countState.write(count);
+ out.output(count);
+ }
+ }
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java
index c55781ff84a8..c67ea3ad3e5b 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/GroupByKeyTranslatorBatch.java
@@ -18,13 +18,13 @@
package org.apache.beam.runners.spark.structuredstreaming.translation.batch;
import static org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY;
+import static org.apache.beam.runners.spark.coders.CoderHelpers.toByteArray;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGlobalGroupBy;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGroupByWindow;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.explodeWindowedKey;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueKey;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueValue;
import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.windowedKV;
-import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers.toByteArray;
import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.collectionEncoder;
import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.encoderOf;
import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.kvEncoder;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/functions/SideInputValues.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/functions/SideInputValues.java
index 23c8d49c3091..0017f95b2a05 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/functions/SideInputValues.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/functions/SideInputValues.java
@@ -34,8 +34,8 @@
import java.util.Map;
import java.util.function.Function;
import javax.annotation.Nullable;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.structuredstreaming.translation.EvaluationContext;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/CoderHelpers.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/CoderHelpers.java
deleted file mode 100644
index f8c63bc34f14..000000000000
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/CoderHelpers.java
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.runners.spark.structuredstreaming.translation.helpers;
-
-import java.io.IOException;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.util.CoderUtils;
-
-/** Serialization utility class. */
-public final class CoderHelpers {
- private CoderHelpers() {}
-
- /**
- * Utility method for serializing an object using the specified coder.
- *
- * @param value Value to serialize.
- * @param coder Coder to serialize with.
- * @param type of value that is serialized
- * @return Byte array representing serialized object.
- */
- public static byte[] toByteArray(T value, Coder coder) {
- try {
- return CoderUtils.encodeToByteArray(coder, value);
- } catch (IOException e) {
- throw new IllegalStateException("Error encoding value: " + value, e);
- }
- }
-
- /**
- * Utility method for deserializing a byte array using the specified coder.
- *
- * @param serialized bytearray to be deserialized.
- * @param coder Coder to deserialize with.
- * @param Type of object to be returned.
- * @return Deserialized object.
- */
- public static T fromByteArray(byte[] serialized, Coder coder) {
- try {
- return CoderUtils.decodeFromByteArray(coder, serialized);
- } catch (IOException e) {
- throw new IllegalStateException("Error decoding bytes for coder: " + coder, e);
- }
- }
-}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
index 29f01c84c02e..08ce47c67bb8 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
@@ -36,6 +36,7 @@
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
From 59a8647c15a669d55a676d6ab94e7b49b791eb6a Mon Sep 17 00:00:00 2001
From: Ryan McGinnis
Date: Mon, 19 Oct 2026 08:52:58 -0700
Subject: [PATCH 038/108] Update nlp_tensorflow_streaming.ipynb (#40641)
* Update nlp_tensorflow_streaming.ipynb
Update link to YouTube Data API. Also, removes link to code samples because they aren't actively being maintained and/or they're in the process of being updated.
* Update examples/notebooks/beam-ml/nlp_tensorflow_streaming.ipynb
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---------
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
examples/notebooks/beam-ml/nlp_tensorflow_streaming.ipynb | 6 ++----
1 file changed, 2 insertions(+), 4 deletions(-)
diff --git a/examples/notebooks/beam-ml/nlp_tensorflow_streaming.ipynb b/examples/notebooks/beam-ml/nlp_tensorflow_streaming.ipynb
index 6f5048e7e8ee..7785041253e6 100644
--- a/examples/notebooks/beam-ml/nlp_tensorflow_streaming.ipynb
+++ b/examples/notebooks/beam-ml/nlp_tensorflow_streaming.ipynb
@@ -54,7 +54,7 @@
"source": [
"Natural Language Processing or NLP is a field of Artifical Intelligence that enables computers to interpret and understand human language. It involves multiple steps such as applying various preprocessing functions, getting predictions from a model, storing the predictions in a useful format, etc.\n",
"Sentiment Analysis is a popular use case of NLP, which allows computers to analyze the sentiment of a text. This notebook demonstrates the use of streaming pipelines in NLP.\n",
- "* Extracts comments using [Youtube API](https://developers.google.com/youtube/v3) and publishing them to Pub/Sub\n",
+ "* Extracts comments using [Youtube API](https://developers.google.com/youtube/data) and publishing them to Pub/Sub\n",
"* Trains a TensorFlow model to predict the sentiment of text\n",
"* Stores the model in Google Cloud and creates a model handler\n",
"* Builds a Beam pipeline to:\n",
@@ -724,9 +724,7 @@
},
"source": [
"## Extracting and sending comments to Pub/Sub\n",
- "YouTube API provides an interface for accessing YouTube data. First, we need to enable YouTube API on the Google Cloud Console project. After that, we need to create a credential, which will further provide an API key. This API key, along with a video ID can be used to access data of that YouTube video. The Publisher created earlier is used to publish each comment to Pub/Sub.\n",
- "\n",
- "See examples of using the YouTube API [here](https://developers.google.com/youtube/v3/code_samples/code_snippets)."
+ "YouTube API provides an interface for accessing YouTube data. First, we need to enable YouTube API on the Google Cloud Console project. After that, we need to create a credential, which will further provide an API key. This API key, along with a video ID can be used to access data of that YouTube video. The Publisher created earlier is used to publish each comment to Pub/Sub.\n"
]
},
{
From 5dc12703dc95fa123456fa843807c174c0843208 Mon Sep 17 00:00:00 2001
From: pablo rodriguez defino
Date: Mon, 19 Oct 2026 11:30:05 -0700
Subject: [PATCH 039/108] [Prism] Keep a stateful key in progress while a split
primary holds it (#40624)
When the SDK accepts a channel split of a bundle of a stateful stage,
stageState.splitBundle removed the key of every residual element from
inprogressKeys and inprogressKeysByBundle, even when elements of the
primary still carry that key. The residual of that key was then
scheduled at once in a second bundle that ran concurrently with the
primary. Prism gives each bundle a copy of the key's state at bundle
start and replaces the key's state with the bundle's copy at commit,
so the last bundle to commit wins. State updates, and the outputs that
depend on them, are lost without any error.
Serializing the processing of a key is the runner's job. The model
defines state per key and window: ReadModifyWriteState "holds a single
value per key and window. The value can be read, modified, and written
back atomically" (beam_runner_api.proto). The stateful processing
post says the runner "cannot parallelize processing of elements for a
particular key+window", and the programming guide says "State is
always fully scoped only to the current processing key". Splitting is
valid SDK behaviour (beam_fn_api.proto), and the Java and Python
harnesses accept channel splits on stateful stages.
splitBundle now collects the keys of the primary elements and releases
a residual's key only if no primary element has it. A held key is
released when the primary bundle is persisted, as before. The handling
of residual timer holds is unchanged.
Reproduction with the Python SDK: all elements share one key, and a
per-element sleep makes Prism split the bundle.
class CountOnOneKey(beam.DoFn):
COUNT = ReadModifyWriteStateSpec('count', VarIntCoder())
DONE = TimerSpec('done', TimeDomain.WATERMARK)
def process(self, element, count=beam.DoFn.StateParam(COUNT),
done=beam.DoFn.TimerParam(DONE)):
time.sleep(0.25)
count.write((count.read() or 0) + 1)
done.set(GlobalWindow().max_timestamp())
@on_timer(DONE)
def report(self, count=beam.DoFn.StateParam(COUNT)):
yield count.read()
options = PipelineOptions(['--runner=PrismRunner',
'--prism_location=',
'--environment_type=LOOPBACK'])
with beam.Pipeline(options=options) as p:
(p | beam.Create(range(40)) | beam.Map(lambda i: ('k', i))
| beam.ParDo(CountOnOneKey()) | beam.Map(print))
A correct runner prints 40. Prism built from master printed 20 in 10
of 10 runs; with this change it printed 40 in 10 of 10 runs. With
--prism_log_level=debug, the same pipeline on master logs one split at
element 20 of 40, and the residual bundle starts before the primary
commits.
TestElementManager_SplitStatefulBundle covers a split inside the only
key, a split inside one of two keys, and a split between keys. The
first two fail without this change.
Not fixed here, follow-ups:
- splitBundle slices the bundle's elements by the data channel index.
A stateful bundle can hold data of one key and timers of another, and
timers are not counted in the data channel, so the slice can
misattribute elements. See #35771.
- startTriggeredBundle, for aggregate stages, does not check
inprogressKeys before it injects a bundle.
---
.../prism/internal/engine/elementmanager.go | 15 +-
.../engine/elementmanager_split_test.go | 194 ++++++++++++++++++
2 files changed, 207 insertions(+), 2 deletions(-)
create mode 100644 sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go
diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
index 80b4a75ac725..2ed73e2c6639 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go
@@ -2318,9 +2318,20 @@ func (ss *stageState) splitBundle(rb RunBundle, firstResidual int, em *ElementMa
es.es = prim
+ // A key stays assigned to this bundle while any primary element still
+ // carries it. Releasing it would let a residual of the same key run in a
+ // second bundle while the primary is still processing it, so both bundles
+ // would start from the same state and the last commit would drop the
+ // other's state changes.
+ primKeys := set[string]{}
+ for _, e := range prim {
+ primKeys.insert(string(e.keyBytes))
+ }
for _, e := range res {
- delete(ss.inprogressKeysByBundle[rb.BundleID], string(e.keyBytes))
- delete(ss.inprogressKeys, string(e.keyBytes))
+ if k := string(e.keyBytes); !primKeys.present(k) {
+ delete(ss.inprogressKeysByBundle[rb.BundleID], k)
+ delete(ss.inprogressKeys, k)
+ }
if e.IsTimer() {
slog.Warn("Unexpected split on a bundle with timers. See https://github.com/apache/beam/issues/35771 for information.")
diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go
new file mode 100644
index 000000000000..d13552b0a363
--- /dev/null
+++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_split_test.go
@@ -0,0 +1,194 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package engine
+
+import (
+ "fmt"
+ "testing"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
+)
+
+// splitInsideKey returns a split point that keeps the first keep elements of
+// key in the primary and returns the remaining elements of key, and every
+// element after them, as the residual.
+func splitInsideKey(key string, keep int) func(t *testing.T, bundleKeys []string) int {
+ return func(t *testing.T, bundleKeys []string) int {
+ t.Helper()
+ for i, k := range bundleKeys {
+ if k == key {
+ return i + keep
+ }
+ }
+ t.Fatalf("key %q not in bundle %v", key, bundleKeys)
+ return 0
+ }
+}
+
+// splitAfterFirstKey returns a split point that keeps exactly the elements of
+// the first key of the bundle in the primary.
+func splitAfterFirstKey(t *testing.T, bundleKeys []string) int {
+ t.Helper()
+ i := 0
+ for i < len(bundleKeys) && bundleKeys[i] == bundleKeys[0] {
+ i++
+ }
+ if i == len(bundleKeys) {
+ t.Fatalf("bundle %v has a single key, want at least two", bundleKeys)
+ }
+ return i
+}
+
+// TestElementManager_SplitStatefulBundle checks that a channel split of a
+// stateful stage's bundle keeps every key held by a primary element assigned
+// to the primary bundle, so no other bundle processes that key concurrently.
+//
+// Prism gives each bundle a copy of the key's state at bundle start and
+// replaces the key's state with the bundle's copy at commit. If two bundles
+// process the same key at the same time, the last commit discards the state
+// changes of the other.
+func TestElementManager_SplitStatefulBundle(t *testing.T) {
+ tests := []struct {
+ name string
+ // keys has one entry per input element, all in the global window.
+ keys []string
+ // firstResidual picks the split point from the bundle's key order,
+ // which depends on map iteration order.
+ firstResidual func(t *testing.T, bundleKeys []string) int
+ }{
+ {
+ name: "split inside the only key",
+ keys: []string{"K", "K", "K"},
+ firstResidual: splitInsideKey("K", 1),
+ }, {
+ name: "split inside one of two keys",
+ keys: []string{"A", "B", "B"},
+ firstResidual: splitInsideKey("B", 1),
+ }, {
+ name: "split between keys",
+ keys: []string{"A", "A", "B", "B"},
+ firstResidual: splitAfterFirstKey,
+ },
+ }
+ for _, test := range tests {
+ t.Run(test.name, func(t *testing.T) {
+ em := NewElementManager(Config{})
+ em.AddStage("impulse", nil, []string{"input"}, nil)
+ em.AddStage("stateful", []string{"input"}, nil, nil)
+ em.StageStateful("stateful", nil)
+ ss := em.stages["stateful"]
+
+ var es []element
+ for i, k := range test.keys {
+ es = append(es, element{
+ window: window.GlobalWindow{},
+ timestamp: mtime.Time(i),
+ pane: typex.NoFiringPane(),
+ elmBytes: []byte(fmt.Sprintf("%s%d", k, i)),
+ keyBytes: []byte(k),
+ sequence: i,
+ })
+ }
+ em.addPending(ss.AddPending(em, es))
+
+ var nextID int
+ genBundID := func() string {
+ defer func() { nextID++ }()
+ return fmt.Sprintf("inst%03d", nextID)
+ }
+ // startBundle builds the next event time bundle, if any, and
+ // returns it with the key of each of its elements.
+ startBundle := func() (RunBundle, []string, bool) {
+ bundID, ok, _, _ := ss.startEventTimeBundle(mtime.MaxTimestamp, genBundID)
+ if !ok {
+ return RunBundle{}, nil, false
+ }
+ var keys []string
+ for _, e := range ss.inprogress[bundID].es {
+ keys = append(keys, string(e.keyBytes))
+ }
+ return RunBundle{StageID: ss.ID, BundleID: bundID, Watermark: mtime.MaxTimestamp}, keys, true
+ }
+
+ primary, bundleKeys, ok := startBundle()
+ if !ok {
+ t.Fatal("startEventTimeBundle() built no bundle, want one")
+ }
+ if got, want := len(bundleKeys), len(test.keys); got != want {
+ t.Fatalf("bundle %v has %v elements, want %v", bundleKeys, got, want)
+ }
+ split := test.firstResidual(t, bundleKeys)
+ wantHeld := set[string]{}
+ for _, k := range bundleKeys[:split] {
+ wantHeld.insert(k)
+ }
+
+ em.ReturnResiduals(primary, split, PColInfo{}, Residuals{})
+
+ checked := set[string]{}
+ for _, k := range bundleKeys {
+ if checked.present(k) {
+ continue
+ }
+ checked.insert(k)
+ want := wantHeld.present(k)
+ if got := ss.inprogressKeys.present(k); got != want {
+ t.Errorf("after split of %v at %v: key %v in inprogressKeys = %v, want %v", bundleKeys, split, k, got, want)
+ }
+ if got := ss.inprogressKeysByBundle[primary.BundleID].present(k); got != want {
+ t.Errorf("after split of %v at %v: key %v in inprogressKeysByBundle[%v] = %v, want %v", bundleKeys, split, k, primary.BundleID, got, want)
+ }
+ }
+
+ // While the primary runs, only residuals of released keys may
+ // be scheduled.
+ processed := split
+ for {
+ rb, keys, ok := startBundle()
+ if !ok {
+ break
+ }
+ for _, k := range keys {
+ if wantHeld.present(k) {
+ t.Fatalf("bundle %v has key %v while primary bundle %v still processes it", rb.BundleID, k, primary.BundleID)
+ }
+ }
+ processed += len(keys)
+ em.PersistBundle(rb, nil, TentativeData{}, PColInfo{}, Residuals{})
+ }
+
+ // Once the primary commits, the residuals of its keys run.
+ em.PersistBundle(primary, nil, TentativeData{}, PColInfo{}, Residuals{})
+ for {
+ rb, keys, ok := startBundle()
+ if !ok {
+ break
+ }
+ processed += len(keys)
+ em.PersistBundle(rb, nil, TentativeData{}, PColInfo{}, Residuals{})
+ }
+
+ if got, want := processed, len(test.keys); got != want {
+ t.Errorf("processed %v elements, want %v", got, want)
+ }
+ if got := em.livePending.Load(); got != 0 {
+ t.Errorf("live pending = %v after all bundles persisted, want 0", got)
+ }
+ })
+ }
+}
From 05c48794a860fea77b640812b2d469da2310e7ed Mon Sep 17 00:00:00 2001
From: Ahmed Abualsaud <65791736+ahmedabu98@users.noreply.github.com>
Date: Mon, 19 Oct 2026 11:36:54 -0700
Subject: [PATCH 040/108] allow dropping pk or watermark column (#40457)
---
.../sdk/io/iceberg/IcebergScanConfig.java | 78 +++++++++++++----
.../cdc/IncrementalChangelogSource.java | 34 +++++++-
...ergCdcReadSchemaTransformProviderTest.java | 7 +-
.../sdk/io/iceberg/IcebergIOReadTest.java | 18 ----
.../sdk/io/iceberg/IcebergScanConfigTest.java | 77 ++++++++++++----
.../cdc/IncrementalChangelogSourceTest.java | 87 ++++++++++++++++++-
6 files changed, 239 insertions(+), 62 deletions(-)
diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java
index bcd574afdaba..6aebf0aa9b26 100644
--- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java
+++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java
@@ -65,6 +65,7 @@
public abstract class IcebergScanConfig implements Serializable {
private transient @MonotonicNonNull Table cachedTable;
private transient org.apache.iceberg.@MonotonicNonNull Schema cachedProjectedSchema;
+ private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRequestedSchema;
private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRequiredSchema;
private transient @MonotonicNonNull Expression cachedFilter;
private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRecordIdSchema;
@@ -136,14 +137,25 @@ static org.apache.iceberg.Schema resolveSchema(
return selectedFields.isEmpty() ? schema : schema.select(selectedFields);
}
- /** Returns the projected Schema after applying column pruning. */
+ /**
+ * Returns the projected Schema after applying column pruning. For CDC reads, this also includes
+ * the primary key and watermark columns, even when they are not requested.
+ */
public org.apache.iceberg.Schema getProjectedSchema() {
if (cachedProjectedSchema == null) {
- cachedProjectedSchema = resolveSchema(getTable().schema(), getKeepFields(), getDropFields());
+ cachedProjectedSchema = withCdcFields(getRequestedSchema());
}
return cachedProjectedSchema;
}
+ /** Returns the Schema of the columns selected by 'keep' or 'drop'. */
+ public org.apache.iceberg.Schema getRequestedSchema() {
+ if (cachedRequestedSchema == null) {
+ cachedRequestedSchema = resolveSchema(getTable().schema(), getKeepFields(), getDropFields());
+ }
+ return cachedRequestedSchema;
+ }
+
/**
* Returns a Schema that includes all the fields required for a successful read. This includes
* explicitly selected fields and fields referenced in the filter statement.
@@ -151,15 +163,54 @@ public org.apache.iceberg.Schema getProjectedSchema() {
public org.apache.iceberg.Schema getRequiredSchema() {
if (cachedRequiredSchema == null) {
cachedRequiredSchema =
- resolveSchema(
- getTable().schema(),
- getKeepFields(),
- getDropFields(),
- FilterUtils.getReferencedFieldNames(getFilterString()));
+ withCdcFields(
+ resolveSchema(
+ getTable().schema(),
+ getKeepFields(),
+ getDropFields(),
+ FilterUtils.getReferencedFieldNames(getFilterString())));
}
return cachedRequiredSchema;
}
+ /**
+ * CDC reads also need the primary key, whose column stats the changelog scan uses to find update
+ * pairs, and the watermark column, which timestamps rows. Adds the top-level columns holding them
+ * if the projection leaves them out.
+ */
+ private org.apache.iceberg.Schema withCdcFields(org.apache.iceberg.Schema projection) {
+ if (!getUseCdc()) {
+ return projection;
+ }
+ org.apache.iceberg.Schema tableSchema = getTable().schema();
+ Set missing = new HashSet<>(tableSchema.identifierFieldIds());
+ @Nullable String watermarkColumn = getWatermarkColumn();
+ @Nullable NestedField watermarkField =
+ watermarkColumn != null ? tableSchema.findField(watermarkColumn) : null;
+ if (watermarkField != null) {
+ missing.add(watermarkField.fieldId());
+ }
+ missing.removeAll(TypeUtil.getProjectedIds(projection));
+ if (missing.isEmpty()) {
+ return projection;
+ }
+
+ Set addedColumnIds = new HashSet<>();
+ for (NestedField column : TypeUtil.select(tableSchema, missing).columns()) {
+ addedColumnIds.add(column.fieldId());
+ }
+ List columns = new ArrayList<>();
+ for (NestedField column : tableSchema.columns()) {
+ @Nullable NestedField projected = projection.asStruct().field(column.fieldId());
+ if (addedColumnIds.contains(column.fieldId())) {
+ columns.add(column);
+ } else if (projected != null) {
+ columns.add(projected);
+ }
+ }
+ return new org.apache.iceberg.Schema(columns);
+ }
+
public org.apache.iceberg.Schema recordIdSchema() {
if (cachedRecordIdSchema == null) {
org.apache.iceberg.Schema fullSchema = getTable().schema();
@@ -438,17 +489,9 @@ void validate(Table table) {
+ invalidOptions));
}
} else {
- Set primaryKeyIds = new HashSet<>(table.schema().identifierFieldIds());
checkState(
- !primaryKeyIds.isEmpty(),
+ !table.schema().identifierFieldIds().isEmpty(),
"Cannot read CDC records as the table schema does not specified any primary key fields.");
- Set projectedFieldIds = TypeUtil.getProjectedIds(getProjectedSchema());
- primaryKeyIds.removeAll(projectedFieldIds);
- checkArgument(
- primaryKeyIds.isEmpty(),
- "When reading CDC records, the projected schema must not drop primary key fields. "
- + "The specified configuration drops the following PK fields: %s",
- primaryKeyIds);
validateMetadataColumns(table);
}
@@ -527,9 +570,6 @@ void validate(Table table) {
error("'watermark_column' must be a timestamp-typed column, but '%s' has type %s"),
watermarkColumn,
field.type().typeId());
- checkArgumentNotNull(
- getProjectedSchema().findField(watermarkColumn),
- "'watermark_column' column should not be dropped.");
}
@Nullable String watermarkColumnTimeUnit = getWatermarkColumnTimeUnit();
diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java
index fe6240d5dc53..9a5d2413717c 100644
--- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java
+++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSource.java
@@ -31,7 +31,9 @@
import org.apache.beam.sdk.io.iceberg.IcebergUtils;
import org.apache.beam.sdk.io.iceberg.ReadUtils;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.transforms.Cast;
import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -167,7 +169,21 @@ public PCollection expand(PBegin input) {
watermarkColumn, scanConfig.getWatermarkColumnTimeUnit())));
}
- return merged.setRowSchema(outputRowSchema);
+ Schema requestedRowSchema =
+ CdcOutputUtils.outputSchema(
+ scanConfig,
+ IcebergUtils.icebergSchemaToBeamSchema(
+ scanConfig.getRequestedSchema(), scanConfig.getUpdateCompatibilityVersion()));
+ if (requestedRowSchema.equals(outputRowSchema)) {
+ return merged.setRowSchema(outputRowSchema);
+ }
+ // the primary key or watermark column was read only for internal use
+ return merged
+ .setRowSchema(outputRowSchema)
+ .apply(
+ "Drop Internal Columns",
+ ParDo.of(new DropInternalColumns(outputRowSchema, requestedRowSchema)))
+ .setRowSchema(requestedRowSchema);
}
/**
@@ -208,4 +224,20 @@ private PCollection boundedSnapshots(PBegin input) {
.collect(Collectors.toList());
return input.apply("Create Snapshot Range", Create.timestamped(timestamped));
}
+
+ /** Drops columns read only for internal use. Rows keep their ValueKind and timestamp. */
+ private static class DropInternalColumns extends DoFn {
+ private final Schema inputSchema;
+ private final Schema outputSchema;
+
+ DropInternalColumns(Schema inputSchema, Schema outputSchema) {
+ this.inputSchema = inputSchema;
+ this.outputSchema = outputSchema;
+ }
+
+ @ProcessElement
+ public void process(@Element Row row, OutputReceiver out) {
+ out.output(Cast.castRow(row, inputSchema, outputSchema));
+ }
+ }
}
diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java
index 3a21566c3056..7a489b59a0b7 100644
--- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java
+++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergCdcReadSchemaTransformProviderTest.java
@@ -230,7 +230,8 @@ public void testManagedReadWithProjectionFilterWatermarkAndSnapshotRange() throw
configMap.put("catalog_properties", properties);
configMap.put("from_snapshot", table.currentSnapshot().snapshotId());
configMap.put("to_snapshot", table.currentSnapshot().snapshotId());
- configMap.put("keep", ImmutableList.of("id", "data", "event_micros"));
+ // leaves out the primary key and the watermark column
+ configMap.put("keep", ImmutableList.of("data"));
configMap.put("filter", "\"category\" = 'include'");
configMap.put("watermark_column", "event_micros");
configMap.put(
@@ -241,7 +242,7 @@ public void testManagedReadWithProjectionFilterWatermarkAndSnapshotRange() throw
IcebergCdcMetadataColumns.ROW_ID,
IcebergCdcMetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER));
- org.apache.iceberg.Schema projectedSchema = table.schema().select("id", "data", "event_micros");
+ org.apache.iceberg.Schema projectedSchema = table.schema().select("data");
Schema recordSchema = IcebergUtils.icebergSchemaToBeamSchema(projectedSchema);
Schema outputSchema =
Schema.builder()
@@ -263,9 +264,7 @@ public void testManagedReadWithProjectionFilterWatermarkAndSnapshotRange() throw
Row record = IcebergUtils.icebergRecordToBeamRow(recordSchema, records.get(i));
return Row.withSchema(outputSchema)
.addValues(
- record.getInt64("id"),
record.getString("data"),
- record.getInt64("event_micros"),
snapshotId,
sequenceNumber,
firstRowId + i,
diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java
index 7920fc37c84b..9a8076bbb7dd 100644
--- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java
+++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergIOReadTest.java
@@ -238,19 +238,6 @@ public void testCdcFailsWhenTableHasNoIdentifierFields() {
read.expand(PBegin.in(testPipeline));
}
- @Test
- public void testCdcFailsWhenProjectionDropsIdentifierFields() {
- assumeTrue(useIncrementalScan);
- TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName());
- warehouse.createTable(tableId, schemaForMode(TestFixtures.SCHEMA, 1));
- IcebergIO.ReadRows read =
- IcebergIO.readRows(catalogConfig()).from(tableId).withCdc().dropping(singletonList("id"));
-
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("projected schema must not drop primary key fields");
- read.expand(PBegin.in(testPipeline));
- }
-
@Test
public void testFailWhenDropAndKeepAreSet() {
TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName());
@@ -410,11 +397,6 @@ public void testScanSelectedFields() throws Exception {
return null;
});
- if (useIncrementalScan) {
- testPipeline.run();
- return;
- }
-
// test drop fields
read = read.keeping(null).dropping(singletonList("id"));
PCollection outputDrop =
diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java
index e928d3cf8f01..989e08dd35d8 100644
--- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java
+++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfigTest.java
@@ -20,6 +20,7 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsString;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertThrows;
import java.util.List;
@@ -76,31 +77,77 @@ public void cdcValidationRequiresIdentifierFields() {
}
@Test
- public void cdcValidationRejectsProjectionDroppingIdentifierFields() {
+ public void cdcProjectionReadsKeyAndWatermarkLeftOutOfKeepOrDrop() {
TableIdentifier tableId = uniqueTableId();
Table table = warehouse.createTable(tableId, CDC_SCHEMA);
- IcebergScanConfig keepWithoutPk =
+ IcebergScanConfig keep =
scanConfigBuilder(tableId, CDC_SCHEMA)
.setUseCdc(true)
.setKeepFields(ImmutableList.of("data"))
+ .setWatermarkColumn("event_time")
+ .setFilterString("\"category\" = 'x'")
.build();
- IllegalArgumentException keepException =
- assertThrows(IllegalArgumentException.class, () -> keepWithoutPk.validate(table));
- assertThat(
- keepException.getMessage(),
- containsString("projected schema must not drop primary key fields"));
+ keep.validate(table);
+ assertEquals(ImmutableList.of("data"), fieldNames(keep.getRequestedSchema()));
+ assertEquals(
+ ImmutableList.of("id", "data", "event_time"), fieldNames(keep.getProjectedSchema()));
+ assertEquals(
+ ImmutableList.of("id", "data", "category", "event_time"),
+ fieldNames(keep.getRequiredSchema()));
+
+ IcebergScanConfig drop =
+ scanConfigBuilder(tableId, CDC_SCHEMA)
+ .setUseCdc(true)
+ .setDropFields(ImmutableList.of("id", "event_time"))
+ .setWatermarkColumn("event_time")
+ .build();
+ drop.validate(table);
+ assertEquals(
+ ImmutableList.of(
+ "data", "category", "event_micros", "optional_time", "required_text", "nested"),
+ fieldNames(drop.getRequestedSchema()));
+ assertEquals(fieldNames(CDC_SCHEMA), fieldNames(drop.getProjectedSchema()));
- IcebergScanConfig dropPk =
+ // a projection that already has both is unchanged
+ IcebergScanConfig unchanged =
scanConfigBuilder(tableId, CDC_SCHEMA)
.setUseCdc(true)
- .setDropFields(ImmutableList.of("id"))
+ .setKeepFields(ImmutableList.of("id", "event_time"))
+ .setWatermarkColumn("event_time")
.build();
- IllegalArgumentException dropException =
- assertThrows(IllegalArgumentException.class, () -> dropPk.validate(table));
- assertThat(
- dropException.getMessage(),
- containsString("projected schema must not drop primary key fields"));
+ assertSame(unchanged.getRequestedSchema(), unchanged.getProjectedSchema());
+
+ // batch reads only project the requested columns
+ IcebergScanConfig batch =
+ scanConfigBuilder(tableId, CDC_SCHEMA).setKeepFields(ImmutableList.of("data")).build();
+ assertEquals(ImmutableList.of("data"), fieldNames(batch.getProjectedSchema()));
+ }
+
+ @Test
+ public void cdcProjectionReadsWholeColumnHoldingNestedKey() {
+ TableIdentifier tableId = uniqueTableId();
+ Table table =
+ warehouse.createTable(
+ tableId,
+ new org.apache.iceberg.Schema(
+ ImmutableList.of(
+ Types.NestedField.required(
+ 1,
+ "key",
+ Types.StructType.of(
+ Types.NestedField.required(2, "id", Types.LongType.get()),
+ Types.NestedField.optional(3, "region", Types.StringType.get()))),
+ Types.NestedField.optional(4, "data", Types.StringType.get())),
+ ImmutableSet.of(2)));
+ IcebergScanConfig scanConfig =
+ scanConfigBuilder(tableId, table.schema())
+ .setUseCdc(true)
+ .setKeepFields(ImmutableList.of("data"))
+ .build();
+
+ // whole columns, as ResolveChanges emits them
+ assertEquals(table.schema().asStruct(), scanConfig.getProjectedSchema().asStruct());
}
@Test
@@ -219,8 +266,6 @@ public void watermarkColumnRejectsInvalidConfigurations() {
true,
ImmutableList.of("id", "required_text"),
"must be a timestamp-typed column");
- assertInvalidWatermark(
- tableId, table, "event_time", true, ImmutableList.of("id"), "should not be dropped");
}
private void assertInvalidWatermark(
diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java
index 27a65a0be8fd..c055e609f765 100644
--- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java
+++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/IncrementalChangelogSourceTest.java
@@ -460,6 +460,79 @@ public void watermarkColumnRestampsProjectedRows() throws Exception {
pipeline.run().waitUntilFinish();
}
+ /**
+ * 'keep' and 'drop' may leave out the primary key and watermark column. The source still reads
+ * them to pair updates and timestamp rows.
+ */
+ @Test
+ public void keyAndWatermarkCanBeLeftOutOfOutput() throws Exception {
+ LocalDateTime t1 =
+ LocalDateTime.ofInstant(
+ java.time.Instant.ofEpochMilli(System.currentTimeMillis() - 3_000L), ZoneOffset.UTC);
+ LocalDateTime t2 = t1.plusSeconds(1);
+ LocalDateTime t3 = t1.plusSeconds(2);
+ // SPLIT_SIZE=1 resolves updates through the CoGroupByKey, the default resolves them locally
+ Map reads =
+ ImmutableMap.of(
+ "Keep",
+ eventChanges("1", t1, t2, t3).setKeepFields(ImmutableList.of("data")).build(),
+ "Drop",
+ eventChanges(String.valueOf(TableProperties.SPLIT_SIZE_DEFAULT), t1, t2, t3)
+ .setDropFields(ImmutableList.of("id", "event_time"))
+ .build());
+
+ for (Map.Entry read : reads.entrySet()) {
+ PCollection rows =
+ pipeline.apply(read.getKey(), new IncrementalChangelogSource(read.getValue()));
+ assertEquals(
+ Schema.builder().addNullableField("data", Schema.FieldType.STRING).build(),
+ rows.getSchema());
+ PAssert.that(
+ rows.apply(read.getKey() + " Format", ParDo.of(new FormatValueKindAndValues()))
+ .apply(read.getKey() + " Timestamps", Reify.timestamps()))
+ .containsInAnyOrder(
+ TimestampedValue.of("INSERT:[a]", millis(t1)),
+ TimestampedValue.of("INSERT:[b]", millis(t2)),
+ // only event_time changed, so pairing compared full rows
+ TimestampedValue.of("UPDATE_BEFORE:[a]", millis(t1)),
+ TimestampedValue.of("UPDATE_AFTER:[a]", millis(t3)),
+ TimestampedValue.of("DELETE:[b]", millis(t2)),
+ TimestampedValue.of("INSERT:[c]", millis(t3)));
+ }
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ /**
+ * Commits id 1 and 2, then an overwrite that changes only id 1's event_time, deletes 2, adds 3.
+ */
+ private IcebergScanConfig.Builder eventChanges(
+ String splitSize, LocalDateTime t1, LocalDateTime t2, LocalDateTime t3) throws IOException {
+ TableIdentifier tableId = TableIdentifier.of("default", testName.getMethodName() + splitSize);
+ Table table =
+ warehouse.createTable(
+ tableId,
+ EVENT_SCHEMA,
+ null,
+ ImmutableMap.of(
+ TableProperties.FORMAT_VERSION, "2", TableProperties.SPLIT_SIZE, splitSize));
+ DataFile v1 =
+ commitAppend(
+ table,
+ splitSize + "-v1.parquet",
+ ImmutableList.of(eventRecord(1L, "a", t1), eventRecord(2L, "b", t2)));
+ commitOverwrite(
+ table, splitSize + "-v2.parquet", v1, eventRecord(1L, "a", t3), eventRecord(3L, "c", t3));
+ return baseConfigBuilder(table, tableId)
+ .setWatermarkColumn("event_time")
+ .setStartingStrategy(StartingStrategy.EARLIEST)
+ .setToSnapshot(table.currentSnapshot().snapshotId());
+ }
+
+ private static Instant millis(LocalDateTime time) {
+ return new Instant(time.toInstant(ZoneOffset.UTC).toEpochMilli());
+ }
+
/**
* Bi-directional rows are windowed per snapshot with zero allowed lateness, so a record that
* reaches the CoGroupByKey after the watermark has passed its snapshot's window is dropped
@@ -516,12 +589,10 @@ private void assertNoRecordsDroppedAcrossSnapshots(boolean streaming) throws Exc
}
private DataFile commitOverwrite(
- Table table, String fileName, DataFile replaced, Record replacement) throws IOException {
+ Table table, String fileName, DataFile replaced, Record... replacements) throws IOException {
DataFile file =
warehouse.writeRecords(
- testName.getMethodName() + "-" + fileName,
- table.schema(),
- ImmutableList.of(replacement));
+ testName.getMethodName() + "-" + fileName, table.schema(), Arrays.asList(replacements));
table.newOverwrite().deleteFile(replaced).addFile(file).commit();
table.refresh();
return file;
@@ -583,6 +654,14 @@ public void process(
}
}
+ private static final class FormatValueKindAndValues extends DoFn {
+ @ProcessElement
+ public void process(
+ @Element Row row, ValueKind valueKind, OutputReceiver outputReceiver) {
+ outputReceiver.output(valueKind.name() + ":" + row.getValues());
+ }
+ }
+
private static final class FormatDefaultedChange extends DoFn {
@ProcessElement
public void process(
From 3bef0c0cb0835fc7eaf43c27e5f4627cfbca1dd3 Mon Sep 17 00:00:00 2001
From: Bruno Volpato
Date: Mon, 19 Oct 2026 14:47:20 -0400
Subject: [PATCH 041/108] [Python] Fail bundles on state write errors (#39992)
* [Python] Fail bundles on state write errors
* [Python] Link state fix in changelog
---
CHANGES.md | 1 +
.../runners/worker/bundle_processor.py | 35 +++---
.../runners/worker/bundle_processor_test.py | 111 ++++++++++++++++++
3 files changed, 133 insertions(+), 14 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 1ec1dbe80864..c49c27c929c7 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -100,6 +100,7 @@
* (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()` instead of silently dropping them ([#40519](https://github.com/apache/beam/issues/40519)).
* (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)).
+* (Python) State write failures reported by portable runners now fail the bundle instead of silently losing state updates ([#39992](https://github.com/apache/beam/pull/39992)).
* (Go) Fixed a data race on the Prism runner's artifact cache map in JobServices ([#32656](https://github.com/apache/beam/issues/32656)).
* (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)).
* (Go) Wrap last dataWriter/timerWriter send errors on Close ([#40405](https://github.com/apache/beam/issues/40405)).
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py
index e62e522c00cf..c9f80ce063fd 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -620,6 +620,17 @@ def __iter__(self) -> Iterator[Any]:
coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type(_ConcatIterable)
+def _await_state_commit(futures) -> None:
+ # Resolve every write before reporting errors, including clears followed by
+ # appends. State futures return error responses rather than raising them.
+ responses = [future.get() for future in futures if future]
+ errors = [
+ response.error for response in responses if response and response.error
+ ]
+ if errors:
+ raise RuntimeError('\n'.join(errors))
+
+
class SynchronousBagRuntimeState(userstate.BagRuntimeState):
def __init__(
self,
@@ -647,15 +658,16 @@ def clear(self) -> None:
self._added_elements = []
def commit(self) -> None:
- to_await = None
+ futures = []
if self._cleared:
- to_await = self._state_handler.clear(self._state_key)
+ futures.append(self._state_handler.clear(self._state_key))
if self._added_elements:
- to_await = self._state_handler.extend(
- self._state_key, self._value_coder.get_impl(), self._added_elements)
- if to_await:
- # To commit, we need to wait on the last state request future to complete.
- to_await.get()
+ futures.append(
+ self._state_handler.extend(
+ self._state_key,
+ self._value_coder.get_impl(),
+ self._added_elements))
+ _await_state_commit(futures)
class SynchronousSetRuntimeState(userstate.SetRuntimeState):
@@ -727,9 +739,7 @@ def commit(self) -> None:
all_futures = self._futures
self._futures = []
- for f in all_futures:
- if f:
- f.get()
+ _await_state_commit(all_futures)
class RangeSet:
@@ -898,10 +908,7 @@ def commit(self) -> None:
self._state_key, self._elem_coder.get_impl(), items_to_add))
self._pending_adds = SortedDict()
- if len(futures):
- # To commit, we need to wait on every state request futures to complete.
- for to_await in futures:
- to_await.get()
+ _await_state_commit(futures)
self._cleared = False
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py
index ad2089c99383..ffe76c1eb819 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py
@@ -31,6 +31,7 @@
from apache_beam.runners.portability.fn_api_runner.worker_handlers import StateServicer
from apache_beam.runners.worker import bundle_processor
from apache_beam.runners.worker import operations
+from apache_beam.runners.worker import sdk_worker
from apache_beam.runners.worker.bundle_processor import BeamTransformFactory
from apache_beam.runners.worker.bundle_processor import BundleProcessor
from apache_beam.runners.worker.bundle_processor import DataInputOperation
@@ -430,6 +431,116 @@ def test_user_modified_sdks_need_to_be_installed_in_runtime_env(self):
"beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0-custom"))
+class RuntimeStateCommitTest(unittest.TestCase):
+ STATE_TYPES = (
+ bundle_processor.SynchronousBagRuntimeState,
+ bundle_processor.SynchronousSetRuntimeState,
+ SynchronousOrderedListRuntimeState)
+
+ def setUp(self):
+ # Keep set compaction deterministic; it is exercised separately below.
+ patcher = mock.patch.object(random, 'random', return_value=0)
+ patcher.start()
+ self.addCleanup(patcher.stop)
+
+ def _create_state(self, state_type, append_errors=('', ), clear_error=''):
+ futures = []
+
+ def response_future(error):
+ future = sdk_worker._Future().set(
+ beam_fn_api_pb2.StateResponse(error=error))
+ future.get = mock.Mock(wraps=future.get)
+ futures.append(future)
+ return future
+
+ append_errors = iter(append_errors)
+ underlying = mock.Mock(spec=sdk_worker.StateHandler)
+ underlying.append_raw.side_effect = (
+ lambda *args: response_future(next(append_errors)))
+ underlying.clear.side_effect = lambda *args: response_future(clear_error)
+ underlying.get_raw.return_value = (b'', None)
+ handler = GlobalCachingStateHandler(StateCache(0), underlying)
+ if state_type is SynchronousOrderedListRuntimeState:
+ key = beam_fn_api_pb2.StateKey(
+ ordered_list_user_state=beam_fn_api_pb2.StateKey.OrderedListUserState(
+ ))
+ else:
+ key = beam_fn_api_pb2.StateKey(
+ bag_user_state=beam_fn_api_pb2.StateKey.BagUserState())
+ return state_type(handler, key, StrUtf8Coder()), futures
+
+ def _add(self, state, value):
+ if isinstance(state, SynchronousOrderedListRuntimeState):
+ state.add((timestamp.Timestamp(1), value))
+ else:
+ state.add(value)
+
+ def test_commit_rejects_failed_append(self):
+ for state_type in self.STATE_TYPES:
+ with self.subTest(state_type=state_type):
+ state, _ = self._create_state(
+ state_type, append_errors=('append failed', ))
+ self._add(state, 'value')
+ with self.assertRaisesRegex(RuntimeError, 'append failed'):
+ state.commit()
+
+ def test_commit_rejects_failed_clear(self):
+ for state_type in self.STATE_TYPES:
+ with self.subTest(state_type=state_type):
+ state, _ = self._create_state(state_type, clear_error='clear failed')
+ state.clear()
+ with self.assertRaisesRegex(RuntimeError, 'clear failed'):
+ state.commit()
+
+ def test_clear_then_append_checks_every_response(self):
+ for state_type in self.STATE_TYPES:
+ for clear_error, append_error in (
+ ('', ''), ('clear failed', ''), ('clear failed', 'append failed')):
+ with self.subTest(state_type=state_type,
+ clear_error=clear_error,
+ append_error=append_error):
+ state, futures = self._create_state(
+ state_type, (append_error, ), clear_error)
+ state.clear()
+ self._add(state, 'replacement')
+ if clear_error:
+ with self.assertRaises(RuntimeError) as raised:
+ state.commit()
+ self.assertEqual(
+ str(raised.exception),
+ '\n'.join(
+ error for error in (clear_error, append_error) if error))
+ else:
+ state.commit()
+ self.assertEqual(len(futures), 2)
+ for future in futures:
+ future.get.assert_called_once()
+
+ def test_commit_rejects_failure_in_earlier_append_chunk(self):
+ state, futures = self._create_state(
+ bundle_processor.SynchronousBagRuntimeState,
+ append_errors=('first append failed', ''))
+ state.add('first')
+ state.add('second')
+ with mock.patch.object(sdk_worker.data_plane,
+ '_DEFAULT_SIZE_FLUSH_THRESHOLD',
+ 1):
+ with self.assertRaisesRegex(RuntimeError, 'first append failed'):
+ state.commit()
+ self.assertEqual(len(futures), 2)
+ for future in futures:
+ future.get.assert_called_once()
+
+ def test_commit_rejects_failed_set_compaction(self):
+ state, futures = self._create_state(
+ bundle_processor.SynchronousSetRuntimeState, clear_error='clear failed')
+ with mock.patch.object(random, 'random', return_value=1):
+ state.add('value')
+ with self.assertRaisesRegex(RuntimeError, 'clear failed'):
+ state.commit()
+ self.assertEqual(len(futures), 2)
+
+
class OrderedListStateTest(unittest.TestCase):
class NoStateCache(StateCache):
def __init__(self):
From 3ab18d061df31be46c619da107e66cfb3ccbd996 Mon Sep 17 00:00:00 2001
From: Elia Liu
Date: Tue, 20 Oct 2026 05:48:44 +1100
Subject: [PATCH 042/108] [Go SDK] Add portable logical type registration and
standard date and micros_instant types (#40154)
* [Go SDK] Add portable logical type registration and standard date and micros_instant types
schema.RegisterLogicalTypeConversion registers a logical type identified
by a URN together with the conversions between its Go type and its
storage type, and derives the row coder from the storage type.
schema.RegisterLogicalTypeCoder does the same for logical types whose
wire encoding differs from the row encoding of the storage type. The
standard beam:logical_type:date:v1 and beam:logical_type:micros_instant:v1
logical types are registered by default as schema.Date and
schema.MicrosInstant.
The standard coder regression test now builds the expected row from the
schema in the coder payload, decodes payload bytes above 0x7f, counts
encoding mismatches as failures, and runs the micros_instant case.
Part of #39684.
* [Go SDK] Replace deprecated reflect.PtrTo in row coder test
* [Go SDK] Handle the zero Date, add MicrosInstant String and Equals, check duplicate registrations
The zero Date now stands for the epoch, MicrosInstant has String and Equals methods, and RegisterLogicalType panics when a Go type or an identifier and argument are already registered for a different logical type.
---
CHANGES.md | 1 +
sdks/go/pkg/beam/core/graph/coder/row.go | 16 +
.../pkg/beam/core/graph/coder/row_decoder.go | 37 ++
.../pkg/beam/core/graph/coder/row_encoder.go | 21 +
sdks/go/pkg/beam/core/graph/coder/row_test.go | 105 +++++
.../runtime/graphx/schema/logicaltypes.go | 142 +++++-
.../core/runtime/graphx/schema/schema_test.go | 30 ++
.../graphx/schema/standard_logical_types.go | 156 +++++++
.../schema/standard_logical_types_test.go | 421 ++++++++++++++++++
.../core/runtime/graphx/serialize_test.go | 16 +
.../regression/coders/fromyaml/fromyaml.go | 210 ++++++---
11 files changed, 1079 insertions(+), 76 deletions(-)
create mode 100644 sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types.go
create mode 100644 sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types_test.go
diff --git a/CHANGES.md b/CHANGES.md
index c49c27c929c7..d50b14cadd1a 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -80,6 +80,7 @@
* (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) that set defaults for all AWS service clients. The existing `--s3_*` options in `S3Options` override these defaults for the S3 client only ([#40528](https://github.com/apache/beam/issues/40528)).
* (Python) Added `--aws_credentials_secret_manager` so AWS credentials (`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and their `--s3_*` overrides) can be secret specs resolved through a secret manager such as `GoogleCloudSecretManager`, instead of plain-text values ([#40528](https://github.com/apache/beam/issues/40528)).
* (Python) Added `AwsSecret` (`AwsSecretsManager` / `type:awssecret`) to retrieve secrets from AWS Secrets Manager, configured from `AwsOptions` (`--aws_*`) with optional per-secret connection overrides ([#40528](https://github.com/apache/beam/issues/40528)).
+* (Go) Added `schema.RegisterLogicalTypeConversion` and `schema.RegisterLogicalTypeCoder` to register portable logical types in one call, and Go types for the standard `beam:logical_type:date:v1` and `beam:logical_type:micros_instant:v1` logical types, so schemas from other SDKs with these types can be used in cross-language pipelines ([#39684](https://github.com/apache/beam/issues/39684)).
## Breaking Changes
diff --git a/sdks/go/pkg/beam/core/graph/coder/row.go b/sdks/go/pkg/beam/core/graph/coder/row.go
index f958d84154c1..0c2a6872dd69 100644
--- a/sdks/go/pkg/beam/core/graph/coder/row.go
+++ b/sdks/go/pkg/beam/core/graph/coder/row.go
@@ -226,3 +226,19 @@ func ReadSimpleRowHeader(fields int, r io.Reader) error {
}
return nil
}
+
+// RowFieldEncoderForType returns an encoding function for values of the given
+// type, using the encoding of that type as a field of a beam schema row.
+//
+// Returns an error if the given type cannot be a field of a beam schema row.
+func RowFieldEncoderForType(rt reflect.Type) (func(any, io.Writer) error, error) {
+ return defaultEnc.fieldEncoderForType(rt)
+}
+
+// RowFieldDecoderForType returns a decoding function for values of the given
+// type, using the encoding of that type as a field of a beam schema row.
+//
+// Returns an error if the given type cannot be a field of a beam schema row.
+func RowFieldDecoderForType(rt reflect.Type) (func(io.Reader) (any, error), error) {
+ return defaultDec.fieldDecoderForType(rt)
+}
diff --git a/sdks/go/pkg/beam/core/graph/coder/row_decoder.go b/sdks/go/pkg/beam/core/graph/coder/row_decoder.go
index 90938c237fe3..2402cafe0e2b 100644
--- a/sdks/go/pkg/beam/core/graph/coder/row_decoder.go
+++ b/sdks/go/pkg/beam/core/graph/coder/row_decoder.go
@@ -95,6 +95,23 @@ func (b *RowDecoderBuilder) decoderForType(t reflect.Type) (func(io.Reader) (any
if t.Kind() == reflect.Ptr {
isPtr = true
t = t.Elem()
+ // A provider for the value type also handles a top level pointer.
+ f, addr, err := b.customFunc(t)
+ if err != nil {
+ return nil, err
+ }
+ if f != nil && !addr {
+ vt := t
+ return func(r io.Reader) (any, error) {
+ v, err := f(r)
+ if err != nil {
+ return nil, err
+ }
+ rv := reflect.New(vt)
+ rv.Elem().Set(reflect.ValueOf(v))
+ return rv.Interface(), nil
+ }, nil
+ }
}
dec, err := b.decoderForStructReflect(t)
if err != nil {
@@ -435,3 +452,23 @@ type typeDecoderFieldReflect struct {
// of the field value (i.e. &foo.bar) and not the field value (i.e. foo.bar).
addr bool
}
+
+// fieldDecoderForType returns a decoder for values of t, using the encoding
+// of t as a row field.
+func (b *RowDecoderBuilder) fieldDecoderForType(t reflect.Type) (func(io.Reader) (any, error), error) {
+ decf, err := b.decoderForSingleTypeReflect(t)
+ if err != nil {
+ return nil, err
+ }
+ return func(r io.Reader) (any, error) {
+ rv := reflect.New(t)
+ v := rv.Elem()
+ if decf.addr {
+ v = rv
+ }
+ if err := decf.decode(v, r); err != nil {
+ return nil, err
+ }
+ return rv.Elem().Interface(), nil
+ }, nil
+}
diff --git a/sdks/go/pkg/beam/core/graph/coder/row_encoder.go b/sdks/go/pkg/beam/core/graph/coder/row_encoder.go
index 6f5c4c04e81c..05d3e85b2d98 100644
--- a/sdks/go/pkg/beam/core/graph/coder/row_encoder.go
+++ b/sdks/go/pkg/beam/core/graph/coder/row_encoder.go
@@ -363,3 +363,24 @@ func (b *RowEncoderBuilder) encoderForStructReflect(t reflect.Type) (func(reflec
return nil
}, nil
}
+
+// fieldEncoderForType returns an encoder for values of t, using the encoding
+// of t as a row field.
+func (b *RowEncoderBuilder) fieldEncoderForType(t reflect.Type) (func(any, io.Writer) error, error) {
+ encf, err := b.encoderForSingleTypeReflect(t)
+ if err != nil {
+ return nil, err
+ }
+ return func(v any, w io.Writer) error {
+ rv := reflect.ValueOf(v)
+ if !rv.IsValid() || (rv.Kind() == reflect.Ptr && rv.IsNil()) {
+ return errors.Errorf("cannot encode a nil value as a row field of type %v", t)
+ }
+ if encf.addr {
+ p := reflect.New(t)
+ p.Elem().Set(rv)
+ rv = p
+ }
+ return encf.encode(rv, w)
+ }, nil
+}
diff --git a/sdks/go/pkg/beam/core/graph/coder/row_test.go b/sdks/go/pkg/beam/core/graph/coder/row_test.go
index c2a925efd306..8b6782457d81 100644
--- a/sdks/go/pkg/beam/core/graph/coder/row_test.go
+++ b/sdks/go/pkg/beam/core/graph/coder/row_test.go
@@ -945,3 +945,108 @@ func TestReflectionRowCoderGeneration_Int16BigEndian(t *testing.T) {
t.Fatalf("dec(enc(%v)) = %v\ndiff (-want, +got): %v", want, got, d)
}
}
+
+func TestRowFieldCoderForType(t *testing.T) {
+ type pair struct {
+ A int64
+ B string
+ }
+ tests := []struct {
+ value any
+ wantBytes []byte
+ }{
+ {value: int64(999), wantBytes: []byte{0xe7, 0x07}},
+ {value: int16(999), wantBytes: []byte{0x03, 0xe7}},
+ {value: "abc", wantBytes: []byte{0x03, 'a', 'b', 'c'}},
+ {value: []byte{1, 2}, wantBytes: []byte{0x02, 1, 2}},
+ {value: true, wantBytes: []byte{0x01}},
+ {value: []int32{1, 2}, wantBytes: []byte{0x00, 0x00, 0x00, 0x02, 0x01, 0x02}},
+ {value: pair{A: 7, B: "x"}, wantBytes: []byte{0x02, 0x00, 0x07, 0x01, 'x'}},
+ }
+ for _, test := range tests {
+ rt := reflect.TypeOf(test.value)
+ t.Run(rt.String(), func(t *testing.T) {
+ enc, err := RowFieldEncoderForType(rt)
+ if err != nil {
+ t.Fatalf("RowFieldEncoderForType(%v) = %v, want nil error", rt, err)
+ }
+ var buf bytes.Buffer
+ if err := enc(test.value, &buf); err != nil {
+ t.Fatalf("enc(%v) = %v, want nil error", test.value, err)
+ }
+ if got := buf.Bytes(); !bytes.Equal(got, test.wantBytes) {
+ t.Fatalf("enc(%v) = %v, want %v", test.value, got, test.wantBytes)
+ }
+ dec, err := RowFieldDecoderForType(rt)
+ if err != nil {
+ t.Fatalf("RowFieldDecoderForType(%v) = %v, want nil error", rt, err)
+ }
+ got, err := dec(bytes.NewBuffer(test.wantBytes))
+ if err != nil {
+ t.Fatalf("dec(%v) = %v, want nil error", test.wantBytes, err)
+ }
+ if d := cmp.Diff(test.value, got); d != "" {
+ t.Fatalf("dec(enc(%v)) = %v\ndiff (-want, +got): %v", test.value, got, d)
+ }
+ })
+ }
+ if _, err := RowFieldEncoderForType(reflect.TypeOf(func() {})); err == nil {
+ t.Errorf("RowFieldEncoderForType(func) = nil error, want error")
+ }
+ enc, err := RowFieldEncoderForType(reflect.TypeOf(int64(0)))
+ if err != nil {
+ t.Fatalf("RowFieldEncoderForType(int64) = %v, want nil error", err)
+ }
+ if err := enc(nil, &bytes.Buffer{}); err == nil {
+ t.Errorf("enc(nil) = nil error, want error")
+ }
+ ptrEnc, err := RowFieldEncoderForType(reflect.TypeOf((*int64)(nil)))
+ if err != nil {
+ t.Fatalf("RowFieldEncoderForType(*int64) = %v, want nil error", err)
+ }
+ if err := ptrEnc((*int64)(nil), &bytes.Buffer{}); err == nil {
+ t.Errorf("enc((*int64)(nil)) = nil error, want error")
+ }
+ if _, err := RowFieldDecoderForType(reflect.TypeOf(func() {})); err == nil {
+ t.Errorf("RowFieldDecoderForType(func) = nil error, want error")
+ }
+}
+
+// A provider for a value type also encodes and decodes a top level pointer to
+// that type.
+func TestRowCoder_TopLevelPointerWithValueProvider(t *testing.T) {
+ type tenths struct{ V int64 }
+ rt := reflect.TypeOf(tenths{})
+ var encB RowEncoderBuilder
+ encB.Register(rt, func(reflect.Type) (func(any, io.Writer) error, error) {
+ return func(v any, w io.Writer) error { return EncodeVarInt(v.(tenths).V*10, w) }, nil
+ })
+ var decB RowDecoderBuilder
+ decB.Register(rt, func(reflect.Type) (func(io.Reader) (any, error), error) {
+ return func(r io.Reader) (any, error) {
+ v, err := DecodeVarInt(r)
+ return tenths{V: v / 10}, err
+ }, nil
+ })
+ prt := reflect.PointerTo(rt)
+ enc, err := encB.Build(prt)
+ if err != nil {
+ t.Fatalf("RowEncoderBuilder.Build(%v) = %v, want nil error", prt, err)
+ }
+ dec, err := decB.Build(prt)
+ if err != nil {
+ t.Fatalf("RowDecoderBuilder.Build(%v) = %v, want nil error", prt, err)
+ }
+ want := &tenths{V: 42}
+ var buf bytes.Buffer
+ if err := enc(want, &buf); err != nil {
+ t.Fatalf("enc(%v) = %v, want nil error", want, err)
+ }
+ got, err := dec(&buf)
+ if err != nil {
+ t.Fatalf("dec = %v, want nil error", err)
+ }
+ if d := cmp.Diff(want, got); d != "" {
+ t.Errorf("dec(enc(%v)) diff (-want, +got): %v", want, d)
+ }
+}
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go
index ee08bbb8f549..bf54fddbc722 100644
--- a/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go
+++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go
@@ -17,9 +17,11 @@ package schema
import (
"fmt"
+ "io"
"reflect"
"sync"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
)
@@ -32,7 +34,9 @@ var (
// A logical type is a type that has distinct representations and storage.
//
// RegisterLogicalType will panic if the storage type of the LogicalType
-// instance is not a valid storage type.
+// instance is not a valid storage type, or if its Go type or its identifier
+// and argument are already registered for a different logical type.
+// Registering the same logical type again has no effect.
func RegisterLogicalType(lt LogicalType) {
defaultRegistry.RegisterLogicalType(lt)
}
@@ -94,7 +98,8 @@ func NewRegistry() *Registry {
}
}
-// RegisterLogicalType a single logical type.
+// RegisterLogicalType registers a single logical type. See the package level
+// RegisterLogicalType for the conditions it panics on.
func (r *Registry) RegisterLogicalType(lt LogicalType) {
// Validates that the storage type has known handling.
st := lt.StorageType()
@@ -113,7 +118,12 @@ func (r *Registry) RegisterLogicalType(lt LogicalType) {
panic(fmt.Sprintf("LogicalType[%v] has the argument %v which does not compare equal to itself", lt.ID(), lt.argV))
}
}
- // TODO add duplication checks.
+ if prev, ok := r.logicalTypeByGoType[lt.GoType()]; ok && !prev.equals(lt) {
+ panic(fmt.Sprintf("Go type %v is already registered as %v -> %v, cannot register it as %v -> %v", lt.GoType(), prev.describe(), prev.StorageType(), lt.describe(), lt.StorageType()))
+ }
+ if prev, ok := r.logicalTypes[lt.key()]; ok && !prev.equals(lt) {
+ panic(fmt.Sprintf("%v is already registered for Go type %v, cannot register it for Go type %v", lt.describe(), prev.GoType(), lt.GoType()))
+ }
r.logicalTypeByGoType[lt.GoType()] = lt
r.logicalTypes[lt.key()] = lt
}
@@ -213,6 +223,21 @@ func (l *LogicalType) key() logicalTypeKey {
return logicalTypeKey{identifier: l.identifier, argument: l.argV.Interface()}
}
+// equals reports whether l and o have the same identifier, argument, Go type
+// and storage type.
+func (l *LogicalType) equals(o LogicalType) bool {
+ return l.goT == o.goT && l.storageT == o.storageT && l.argT == o.argT && l.key() == o.key()
+}
+
+// describe returns the identifier of the logical type and its argument, if
+// any, for messages.
+func (l *LogicalType) describe() string {
+ if l.argT == nil {
+ return fmt.Sprintf("LogicalType[%v]", l.identifier)
+ }
+ return fmt.Sprintf("LogicalType[%v, %v]", l.identifier, l.argV.Interface())
+}
+
func preRegLogicalTypes(r *Registry) {
r.RegisterLogicalType(ToLogicalType("int", reflectx.Int, reflectx.Int64))
r.RegisterLogicalType(ToLogicalType("int8", reflectx.Int8, reflectx.Int64))
@@ -224,4 +249,115 @@ func preRegLogicalTypes(r *Registry) {
func init() {
preRegLogicalTypes(defaultRegistry)
+ registerStandardLogicalTypes(defaultRegistry)
+}
+
+// RegisterLogicalTypeConversion registers a portable logical type identified
+// by urn. Values of GoT are stored in schema rows as values of StorageT,
+// converted with toStorage and fromStorage, and encoded with the row field
+// encoding of StorageT. StorageT must have a schema representation: an atomic
+// type, a struct of such types, or a registered logical type. GoT must be a
+// concrete type. A struct kind GoT with a non row StorageT can be used as a
+// field of a row, but not as a top level schema type.
+//
+// RegisterLogicalTypeConversion must be called before beam.Init(), and
+// conventionally is called in a package init() function.
+func RegisterLogicalTypeConversion[GoT, StorageT any](urn string, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) {
+ registerLogicalTypeConversion(defaultRegistry, ToLogicalType(urn, typeOf[GoT](), typeOf[StorageT]()), toStorage, fromStorage)
+}
+
+// RegisterParameterizedLogicalTypeConversion is RegisterLogicalTypeConversion
+// for a logical type with an argument, such as a timestamp and its precision.
+// Each argument value is registered with its own Go type, and a schema field
+// with the urn and that argument maps to GoT. See ToLogicalTypeWithArgument
+// for the supported argument types.
+func RegisterParameterizedLogicalTypeConversion[GoT, StorageT any](urn string, argument any, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) {
+ registerLogicalTypeConversion(defaultRegistry, ToLogicalTypeWithArgument(urn, typeOf[GoT](), typeOf[StorageT](), argument), toStorage, fromStorage)
+}
+
+// RegisterLogicalTypeCoder registers a portable logical type identified by
+// urn whose wire encoding differs from the row field encoding of StorageT.
+// StorageT is only used as the schema representation of the logical type.
+// enc and dec produce and consume the encoded bytes of a value directly, in
+// the wire encoding that the urn defines. GoT must be a concrete type.
+// StorageT cannot be inferred, so both type parameters are given explicitly,
+// for example
+// RegisterLogicalTypeCoder[Money, []byte]("example:money:v1", encodeMoney, decodeMoney).
+//
+// RegisterLogicalTypeCoder must be called before beam.Init(), and
+// conventionally is called in a package init() function.
+func RegisterLogicalTypeCoder[GoT, StorageT any](urn string, enc func(GoT, io.Writer) error, dec func(io.Reader) (GoT, error)) {
+ registerLogicalTypeCoder(defaultRegistry, ToLogicalType(urn, typeOf[GoT](), typeOf[StorageT]()), enc, dec)
+}
+
+func typeOf[T any]() reflect.Type {
+ return reflect.TypeOf((*T)(nil)).Elem()
+}
+
+// registerLogicalTypeConversion registers lt with the registry and registers
+// row coder providers for its Go type, built from the row field coder of the
+// storage type and the conversion functions.
+func registerLogicalTypeConversion[GoT, StorageT any](r *Registry, lt LogicalType, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) {
+ st := lt.StorageType()
+ checkLogicalTypeKinds(lt)
+ r.RegisterLogicalType(lt)
+ coder.RegisterSchemaProviders(lt.GoType(),
+ func(reflect.Type) (func(any, io.Writer) error, error) {
+ enc, err := coder.RowFieldEncoderForType(st)
+ if err != nil {
+ return nil, err
+ }
+ return func(v any, w io.Writer) error {
+ s, err := toStorage(v.(GoT))
+ if err != nil {
+ return err
+ }
+ return enc(s, w)
+ }, nil
+ },
+ func(reflect.Type) (func(io.Reader) (any, error), error) {
+ dec, err := coder.RowFieldDecoderForType(st)
+ if err != nil {
+ return nil, err
+ }
+ return func(rd io.Reader) (any, error) {
+ s, err := dec(rd)
+ if err != nil {
+ return nil, err
+ }
+ return fromStorage(s.(StorageT))
+ }, nil
+ })
+}
+
+// registerLogicalTypeCoder registers lt with the registry and registers enc
+// and dec as the row coder providers for its Go type.
+func registerLogicalTypeCoder[GoT any](r *Registry, lt LogicalType, enc func(GoT, io.Writer) error, dec func(io.Reader) (GoT, error)) {
+ checkLogicalTypeKinds(lt)
+ r.RegisterLogicalType(lt)
+ coder.RegisterSchemaProviders(lt.GoType(),
+ func(reflect.Type) (func(any, io.Writer) error, error) {
+ return func(v any, w io.Writer) error {
+ return enc(v.(GoT), w)
+ }, nil
+ },
+ func(reflect.Type) (func(io.Reader) (any, error), error) {
+ return func(rd io.Reader) (any, error) {
+ return dec(rd)
+ }, nil
+ })
+}
+
+// checkLogicalTypeKinds panics when the Go type of lt is an interface or its
+// storage type is a pointer. The row coder matches an interface by all of its
+// implementations, and the schema matches only the interface itself. The row
+// coders decide nullability from the Go type of a field, so a pointer storage
+// type behind a non pointer Go type would fail at encoding time.
+func checkLogicalTypeKinds(lt LogicalType) {
+ if lt.GoType().Kind() == reflect.Interface {
+ panic(fmt.Sprintf("LogicalType[%v] has the interface Go type %v, Go types must be concrete", lt.ID(), lt.GoType()))
+ }
+ if lt.StorageType().Kind() == reflect.Ptr {
+ panic(fmt.Sprintf("LogicalType[%v] has the pointer storage type %v, storage types must not be pointers", lt.ID(), lt.StorageType()))
+ }
}
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go
index 08916e3898d0..09caa1417e74 100644
--- a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go
+++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go
@@ -1163,3 +1163,33 @@ func TestRegisterLogicalType_Validation(t *testing.T) {
})
}
}
+
+func TestRegisterLogicalType_Duplicate(t *testing.T) {
+ reg := NewRegistry()
+ lt := ToLogicalType("schema.anyString", anyStringType, reflectx.String)
+ reg.RegisterLogicalType(lt)
+ // Registering the same logical type again has no effect.
+ reg.RegisterLogicalType(lt)
+
+ for name, dup := range map[string]LogicalType{
+ "identifier for another Go type": ToLogicalType("schema.anyString", parameterizedStringType, reflectx.String),
+ "Go type with another identifier": ToLogicalType("schema.otherString", anyStringType, reflectx.String),
+ "Go type with another storage type": ToLogicalType("schema.anyString", anyStringType, reflectx.ByteSlice),
+ "Go type with an argument": ToLogicalTypeWithArgument("schema.anyString", anyStringType, reflectx.String, int32(5)),
+ } {
+ t.Run(name, func(t *testing.T) {
+ defer func() {
+ if recover() == nil {
+ t.Errorf("RegisterLogicalType did not panic for a duplicate %v", name)
+ }
+ }()
+ reg.RegisterLogicalType(dup)
+ })
+ }
+ if got, ok := reg.lookupLogicalType("schema.anyString", reflect.Value{}); !ok || got.GoType() != anyStringType {
+ t.Errorf("lookupLogicalType(schema.anyString) = %v, %v, want %v, true", got.GoType(), ok, anyStringType)
+ }
+ if _, ok := reg.logicalTypeByGoType[parameterizedStringType]; ok {
+ t.Errorf("logicalTypeByGoType has %v after a rejected registration", parameterizedStringType)
+ }
+}
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types.go
new file mode 100644
index 000000000000..dcc12a45e69f
--- /dev/null
+++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types.go
@@ -0,0 +1,156 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package schema
+
+import (
+ "time"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+)
+
+// URNs of the standard logical types defined in schema.proto.
+const (
+ // URNDate identifies the Date logical type.
+ URNDate = "beam:logical_type:date:v1"
+ // URNMicrosInstant identifies the MicrosInstant logical type.
+ URNMicrosInstant = "beam:logical_type:micros_instant:v1"
+)
+
+// Date is a calendar date without a time zone or time of day, the
+// beam:logical_type:date:v1 logical type. It is stored as the number of days
+// since 1970-01-01 in an INT64 field. The zero Date is the epoch, 1970-01-01.
+// Time and String normalize out of range fields as time.Date does. Encoding
+// fails unless the fields form a valid date with a year from -999999999 to
+// 999999999, and decoding fails for a day count outside that range.
+type Date struct {
+ Year int
+ Month time.Month
+ Day int
+}
+
+// DateOf returns the Date of t in t's location.
+func DateOf(t time.Time) Date {
+ y, m, d := t.Date()
+ return Date{Year: y, Month: m, Day: d}
+}
+
+// epochDate is the date that the zero Date stands for.
+var epochDate = Date{Year: 1970, Month: time.January, Day: 1}
+
+// Time returns the start of the date in UTC.
+func (d Date) Time() time.Time {
+ if d == (Date{}) {
+ d = epochDate
+ }
+ return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, time.UTC)
+}
+
+// String returns the date in ISO 8601 format, such as 2006-01-02.
+func (d Date) String() string {
+ return d.Time().Format("2006-01-02")
+}
+
+const secondsPerDay = 24 * 60 * 60
+
+// minDateDays and maxDateDays bound the supported days since the epoch.
+var (
+ minDateDays = time.Date(-999999999, time.January, 1, 0, 0, 0, 0, time.UTC).Unix() / secondsPerDay
+ maxDateDays = time.Date(999999999, time.December, 31, 0, 0, 0, 0, time.UTC).Unix() / secondsPerDay
+)
+
+func dateToStorage(d Date) (int64, error) {
+ if d == (Date{}) {
+ d = epochDate
+ }
+ t := d.Time()
+ // Out of range fields can overflow in time.Date, so only a date that
+ // round trips is encoded.
+ if DateOf(t) != d {
+ return 0, errors.Errorf("%d-%02d-%02d is not a valid date", d.Year, d.Month, d.Day)
+ }
+ days := t.Unix() / secondsPerDay
+ if days < minDateDays || days > maxDateDays {
+ return 0, errors.Errorf("date %v is out of range", d)
+ }
+ return days, nil
+}
+
+func dateFromStorage(days int64) (Date, error) {
+ if days < minDateDays || days > maxDateDays {
+ return Date{}, errors.Errorf("%v days since the epoch is out of the supported date range", days)
+ }
+ return DateOf(time.Unix(days*secondsPerDay, 0).UTC()), nil
+}
+
+// MicrosInstant is a timestamp with microsecond precision, the
+// beam:logical_type:micros_instant:v1 logical type. It is stored as a row of
+// the seconds and the microseconds since the epoch. Encoding a value with
+// sub microsecond precision fails; truncate such values with time.Time.Truncate
+// first.
+type MicrosInstant time.Time
+
+// Time returns the instant as a time.Time.
+func (m MicrosInstant) Time() time.Time {
+ return time.Time(m)
+}
+
+// String returns the instant in RFC 3339 format with microseconds, such as
+// 2020-08-13T14:14:14.123456Z.
+func (m MicrosInstant) String() string {
+ return time.Time(m).Format("2006-01-02T15:04:05.000000Z07:00")
+}
+
+// Equals reports whether m and o are the same instant.
+func (m MicrosInstant) Equals(o MicrosInstant) bool {
+ return time.Time(m).Equal(time.Time(o))
+}
+
+// microsInstantStorage is the storage type of MicrosInstant.
+type microsInstantStorage struct {
+ Seconds int64 `beam:"seconds"`
+ Micros int64 `beam:"micros"`
+}
+
+func microsInstantToStorage(m MicrosInstant) (microsInstantStorage, error) {
+ t := time.Time(m)
+ nanos := t.Nanosecond()
+ if nanos%1000 != 0 {
+ return microsInstantStorage{}, errors.Errorf("MicrosInstant %v has sub microsecond precision", t)
+ }
+ return microsInstantStorage{Seconds: t.Unix(), Micros: int64(nanos / 1000)}, nil
+}
+
+func microsInstantFromStorage(s microsInstantStorage) (MicrosInstant, error) {
+ if s.Micros < 0 || s.Micros > 999999 {
+ return MicrosInstant{}, errors.Errorf("MicrosInstant micros %v is out of range [0, 999999]", s.Micros)
+ }
+ return MicrosInstant(time.Unix(s.Seconds, s.Micros*1000).UTC()), nil
+}
+
+// registerStandardLogicalTypes registers the standard logical types with the
+// registry.
+func registerStandardLogicalTypes(r *Registry) {
+ registerLogicalTypeConversion(r, ToLogicalType(URNDate, typeOf[Date](), reflectx.Int64), dateToStorage, dateFromStorage)
+ registerLogicalTypeConversion(r, ToLogicalType(URNMicrosInstant, typeOf[MicrosInstant](), typeOf[microsInstantStorage]()), microsInstantToStorage, microsInstantFromStorage)
+}
+
+func init() {
+ // Pipeline graphs serialize the standard logical types by name.
+ runtime.RegisterType(typeOf[Date]())
+ runtime.RegisterType(typeOf[MicrosInstant]())
+}
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types_test.go
new file mode 100644
index 000000000000..88b86be1248e
--- /dev/null
+++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types_test.go
@@ -0,0 +1,421 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package schema
+
+import (
+ "bytes"
+ "encoding/binary"
+ "io"
+ "math"
+ "reflect"
+ "testing"
+ "time"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
+ pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
+ "github.com/google/go-cmp/cmp"
+ "google.golang.org/protobuf/encoding/prototext"
+ "google.golang.org/protobuf/testing/protocmp"
+)
+
+func atomicType(at pipepb.AtomicType) *pipepb.FieldType {
+ return &pipepb.FieldType{TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: at}}
+}
+
+func logicalFieldType(urn string, representation *pipepb.FieldType) *pipepb.FieldType {
+ return &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_LogicalType{
+ LogicalType: &pipepb.LogicalType{
+ Urn: urn,
+ Representation: representation,
+ },
+ },
+ }
+}
+
+// standardRow uses the standard logical types as fields, with the same field
+// names as the micros_instant case of standard_coders.yaml.
+type standardRow struct {
+ F_timestamp MicrosInstant `beam:"f_timestamp"`
+ F_string string `beam:"f_string"`
+ F_int int64 `beam:"f_int"`
+ F_date Date `beam:"f_date"`
+}
+
+func TestStandardLogicalTypes_Schema(t *testing.T) {
+ rt := reflect.TypeOf(standardRow{})
+ want := &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ {Name: "f_timestamp", Type: logicalFieldType(URNMicrosInstant, &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_RowType{RowType: &pipepb.RowType{Schema: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ {Name: "seconds", Type: atomicType(pipepb.AtomicType_INT64)},
+ {Name: "micros", Type: atomicType(pipepb.AtomicType_INT64)},
+ },
+ }}},
+ })},
+ {Name: "f_string", Type: atomicType(pipepb.AtomicType_STRING)},
+ {Name: "f_int", Type: atomicType(pipepb.AtomicType_INT64)},
+ {Name: "f_date", Type: logicalFieldType(URNDate, atomicType(pipepb.AtomicType_INT64))},
+ },
+ }
+ got, err := FromType(rt)
+ if err != nil {
+ t.Fatalf("FromType(%v) = %v, want nil error", rt, err)
+ }
+ if d := cmp.Diff(want, got, protocmp.Transform(), protocmp.IgnoreFields(&pipepb.Schema{}, "id")); d != "" {
+ t.Errorf("FromType(%v) diff (-want, +got): %v", rt, d)
+ }
+ // A schema from another SDK has a different id, and must still map to the
+ // standard Go types.
+ want.Id = "standard-logical-types-test"
+ gotType, err := ToType(want)
+ if err != nil {
+ t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(want), err)
+ }
+ for i, wantField := range []reflect.Type{typeOf[MicrosInstant](), typeOf[string](), typeOf[int64](), typeOf[Date]()} {
+ if got := gotType.Field(i).Type; got != wantField {
+ t.Errorf("ToType field %d type = %v, want %v", i, got, wantField)
+ }
+ }
+}
+
+func TestStandardLogicalTypes_RowEncoding(t *testing.T) {
+ // The bytes are the micros_instant example of standard_coders.yaml, with a
+ // trailing date field appended: 2020-08-13 is 18487 days since the epoch.
+ want := standardRow{
+ F_timestamp: MicrosInstant(time.Date(2020, 8, 13, 14, 14, 14, 123456000, time.UTC)),
+ F_string: "2020-08-13T14:14:14.123456Z",
+ F_int: 1597328054123456,
+ F_date: Date{Year: 2020, Month: time.August, Day: 13},
+ }
+ wantBytes := []byte("\x04\x00\x02\x00\xb6\x95\xd5\xf9\x05\xc0\xc4\x07\x1b2020-08-13T14:14:14.123456Z\xc0\xf7\x85\xda\xae\x98\xeb\x02\xb7\x90\x01")
+ rt := reflect.TypeOf(want)
+ enc, err := coder.RowEncoderForStruct(rt)
+ if err != nil {
+ t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err)
+ }
+ var buf bytes.Buffer
+ if err := enc(want, &buf); err != nil {
+ t.Fatalf("enc(%v) = %v, want nil error", want, err)
+ }
+ if got := buf.Bytes(); !bytes.Equal(got, wantBytes) {
+ t.Fatalf("enc(%v) = %q, want %q", want, got, wantBytes)
+ }
+ dec, err := coder.RowDecoderForStruct(rt)
+ if err != nil {
+ t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err)
+ }
+ got, err := dec(bytes.NewBuffer(wantBytes))
+ if err != nil {
+ t.Fatalf("dec(%q) = %v, want nil error", wantBytes, err)
+ }
+ if d := cmp.Diff(want, got, cmp.Comparer(MicrosInstant.Equals)); d != "" {
+ t.Errorf("dec(enc(%v)) diff (-want, +got): %v", want, d)
+ }
+}
+
+func TestMicrosInstant_SubMicrosecondPrecision(t *testing.T) {
+ enc, err := coder.RowEncoderForStruct(reflect.TypeOf(standardRow{}))
+ if err != nil {
+ t.Fatalf("RowEncoderForStruct = %v, want nil error", err)
+ }
+ row := standardRow{F_timestamp: MicrosInstant(time.Unix(0, 1500))}
+ if err := enc(row, &bytes.Buffer{}); err == nil {
+ t.Errorf("enc(%v) = nil error, want error for sub microsecond precision", row)
+ }
+}
+
+func TestDate_Storage(t *testing.T) {
+ tests := []struct {
+ date Date
+ days int64
+ }{
+ {Date{1970, time.January, 1}, 0},
+ {Date{1969, time.December, 31}, -1},
+ {Date{2020, time.August, 13}, 18487},
+ {Date{1600, time.February, 29}, -135081},
+ {Date{9999, time.December, 31}, 2932896},
+ }
+ for _, test := range tests {
+ days, err := dateToStorage(test.date)
+ if err != nil {
+ t.Fatalf("dateToStorage(%v) = %v, want nil error", test.date, err)
+ }
+ if days != test.days {
+ t.Errorf("dateToStorage(%v) = %v, want %v", test.date, days, test.days)
+ }
+ date, err := dateFromStorage(test.days)
+ if err != nil {
+ t.Fatalf("dateFromStorage(%v) = %v, want nil error", test.days, err)
+ }
+ if date != test.date {
+ t.Errorf("dateFromStorage(%v) = %v, want %v", test.days, date, test.date)
+ }
+ }
+ if got, want := DateOf(time.Date(2020, 8, 13, 23, 59, 59, 0, time.UTC)).String(), "2020-08-13"; got != want {
+ t.Errorf("DateOf(...).String() = %v, want %v", got, want)
+ }
+}
+
+func TestRegisterLogicalTypeConversion(t *testing.T) {
+ type celsius float64
+ type reading struct {
+ Temperature celsius
+ }
+ // Registering with the default registry mirrors what users do in init().
+ RegisterLogicalTypeConversion[celsius, float64]("schema.test:celsius:v1",
+ func(c celsius) (float64, error) { return float64(c) + 273.15, nil },
+ func(k float64) (celsius, error) { return celsius(k - 273.15), nil })
+
+ rt := reflect.TypeOf(reading{})
+ schm, err := FromType(rt)
+ if err != nil {
+ t.Fatalf("FromType(%v) = %v, want nil error", rt, err)
+ }
+ wantField := logicalFieldType("schema.test:celsius:v1", atomicType(pipepb.AtomicType_DOUBLE))
+ if d := cmp.Diff(wantField, schm.GetFields()[0].GetType(), protocmp.Transform()); d != "" {
+ t.Errorf("FromType(%v) field diff (-want, +got): %v", rt, d)
+ }
+
+ enc, err := coder.RowEncoderForStruct(rt)
+ if err != nil {
+ t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err)
+ }
+ var buf bytes.Buffer
+ if err := enc(reading{Temperature: 10}, &buf); err != nil {
+ t.Fatalf("enc = %v, want nil error", err)
+ }
+ // 283.15 as a big endian double, after the row header.
+ wantBytes := []byte{0x01, 0x00, 0x40, 0x71, 0xb2, 0x66, 0x66, 0x66, 0x66, 0x66}
+ if got := buf.Bytes(); !bytes.Equal(got, wantBytes) {
+ t.Fatalf("enc(reading{10}) = %v, want %v", got, wantBytes)
+ }
+ dec, err := coder.RowDecoderForStruct(rt)
+ if err != nil {
+ t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err)
+ }
+ got, err := dec(&buf)
+ if err != nil {
+ t.Fatalf("dec = %v, want nil error", err)
+ }
+ if d := cmp.Diff(reading{Temperature: 10}, got, cmp.Comparer(func(a, b celsius) bool { return a-b < 1e-9 && b-a < 1e-9 })); d != "" {
+ t.Errorf("dec(enc(reading{10})) diff (-want, +got): %v", d)
+ }
+}
+
+// fixedCode is a user logical type with an argument, registered with the
+// exported parameterized registration.
+type fixedCode string
+
+func TestRegisterParameterizedLogicalTypeConversion(t *testing.T) {
+ RegisterParameterizedLogicalTypeConversion[fixedCode, string]("schema.test:fixed_code:v1", int32(4),
+ func(c fixedCode) (string, error) { return string(c), nil },
+ func(s string) (fixedCode, error) { return fixedCode(s), nil })
+ rt := reflect.TypeOf(struct{ Code fixedCode }{})
+ schm, err := FromType(rt)
+ if err != nil {
+ t.Fatalf("FromType(%v) = %v, want nil error", rt, err)
+ }
+ want := logicalFieldType("schema.test:fixed_code:v1", atomicType(pipepb.AtomicType_STRING))
+ want.GetLogicalType().ArgumentType = int32FieldType()
+ want.GetLogicalType().Argument = int32FieldValue(4)
+ if d := cmp.Diff(want, schm.GetFields()[0].GetType(), protocmp.Transform()); d != "" {
+ t.Errorf("FromType(%v) field diff (-want, +got): %v", rt, d)
+ }
+ schm.Id = "fixed-code-from-elsewhere"
+ got, err := ToType(schm)
+ if err != nil {
+ t.Fatalf("ToType = %v, want nil error", err)
+ }
+ if ft := got.Field(0).Type; ft != typeOf[fixedCode]() {
+ t.Errorf("ToType field type = %v, want %v", ft, typeOf[fixedCode]())
+ }
+ enc, err := coder.RowEncoderForStruct(rt)
+ if err != nil {
+ t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err)
+ }
+ var buf bytes.Buffer
+ if err := enc(struct{ Code fixedCode }{"ABCD"}, &buf); err != nil {
+ t.Fatalf("enc = %v, want nil error", err)
+ }
+ if got, want := buf.Bytes(), []byte("\x01\x00\x04ABCD"); !bytes.Equal(got, want) {
+ t.Errorf("enc = %q, want %q", got, want)
+ }
+}
+
+// money is a user logical type with its own wire encoding: the cents as a
+// big endian int32 inside a BYTES representation.
+type money int32
+
+func TestRegisterLogicalTypeCoder(t *testing.T) {
+ RegisterLogicalTypeCoder[money, []byte]("schema.test:money:v1",
+ func(m money, w io.Writer) error {
+ var b [4]byte
+ binary.BigEndian.PutUint32(b[:], uint32(m))
+ return coder.EncodeBytes(b[:], w)
+ },
+ func(r io.Reader) (money, error) {
+ b, err := coder.DecodeBytes(r)
+ if err != nil {
+ return 0, err
+ }
+ return money(int32(binary.BigEndian.Uint32(b))), nil
+ })
+ rt := reflect.TypeOf(struct{ Price money }{})
+ schm, err := FromType(rt)
+ if err != nil {
+ t.Fatalf("FromType(%v) = %v, want nil error", rt, err)
+ }
+ want := logicalFieldType("schema.test:money:v1", atomicType(pipepb.AtomicType_BYTES))
+ if d := cmp.Diff(want, schm.GetFields()[0].GetType(), protocmp.Transform()); d != "" {
+ t.Errorf("FromType(%v) field diff (-want, +got): %v", rt, d)
+ }
+ enc, err := coder.RowEncoderForStruct(rt)
+ if err != nil {
+ t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err)
+ }
+ var buf bytes.Buffer
+ if err := enc(struct{ Price money }{-2}, &buf); err != nil {
+ t.Fatalf("enc = %v, want nil error", err)
+ }
+ wantBytes := []byte{0x01, 0x00, 0x04, 0xff, 0xff, 0xff, 0xfe}
+ if got := buf.Bytes(); !bytes.Equal(got, wantBytes) {
+ t.Fatalf("enc = %v, want %v", got, wantBytes)
+ }
+ dec, err := coder.RowDecoderForStruct(rt)
+ if err != nil {
+ t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err)
+ }
+ got, err := dec(&buf)
+ if err != nil {
+ t.Fatalf("dec = %v, want nil error", err)
+ }
+ if d := cmp.Diff(struct{ Price money }{-2}, got); d != "" {
+ t.Errorf("dec(enc) diff (-want, +got): %v", d)
+ }
+}
+
+func TestRegisterLogicalTypeConversion_PointerStorage(t *testing.T) {
+ type celsiusPtr float64
+ defer func() {
+ if recover() == nil {
+ t.Errorf("RegisterLogicalTypeConversion with a pointer storage type did not panic")
+ }
+ }()
+ RegisterLogicalTypeConversion[celsiusPtr, *float64]("schema.test:celsius_ptr:v1",
+ func(c celsiusPtr) (*float64, error) { f := float64(c); return &f, nil },
+ func(f *float64) (celsiusPtr, error) { return celsiusPtr(*f), nil })
+}
+
+func TestRegisterLogicalType_InterfaceGoType(t *testing.T) {
+ type textual interface{ Text() string }
+ t.Run("conversion", func(t *testing.T) {
+ defer func() {
+ if recover() == nil {
+ t.Errorf("RegisterLogicalTypeConversion with an interface Go type did not panic")
+ }
+ }()
+ RegisterLogicalTypeConversion[textual, string]("schema.test:textual:v1",
+ func(v textual) (string, error) { return v.Text(), nil },
+ func(s string) (textual, error) { return nil, nil })
+ })
+ t.Run("coder", func(t *testing.T) {
+ defer func() {
+ if recover() == nil {
+ t.Errorf("RegisterLogicalTypeCoder with an interface Go type did not panic")
+ }
+ }()
+ RegisterLogicalTypeCoder[textual, string]("schema.test:textual_coder:v1",
+ func(v textual, w io.Writer) error { return nil },
+ func(r io.Reader) (textual, error) { return nil, nil })
+ })
+}
+
+func TestDate_StorageRange(t *testing.T) {
+ for _, days := range []int64{minDateDays, maxDateDays} {
+ date, err := dateFromStorage(days)
+ if err != nil {
+ t.Fatalf("dateFromStorage(%v) = %v, want nil error", days, err)
+ }
+ if got, err := dateToStorage(date); err != nil || got != days {
+ t.Errorf("dateToStorage(%v) = %v, %v, want %v, nil", date, got, err, days)
+ }
+ }
+ for _, days := range []int64{minDateDays - 1, maxDateDays + 1, 1 << 62, math.MinInt64} {
+ if date, err := dateFromStorage(days); err == nil {
+ t.Errorf("dateFromStorage(%v) = %v, want error", days, date)
+ }
+ }
+ for _, date := range []Date{
+ {-1000000000, time.December, 31},
+ {1000000000, time.January, 1},
+ {2020, time.February, 30},
+ {1970, time.January, math.MaxInt/4 + 2},
+ {math.MaxInt / 2, time.January, 1},
+ } {
+ if days, err := dateToStorage(date); err == nil {
+ t.Errorf("dateToStorage(%d-%02d-%02d) = %v, want error", date.Year, date.Month, date.Day, days)
+ }
+ }
+}
+
+func TestMicrosInstant_StorageRange(t *testing.T) {
+ for _, micros := range []int64{-1, 1000000, 1 << 61} {
+ s := microsInstantStorage{Seconds: 0, Micros: micros}
+ if m, err := microsInstantFromStorage(s); err == nil {
+ t.Errorf("microsInstantFromStorage(%+v) = %v, want error", s, m.Time())
+ }
+ }
+ s := microsInstantStorage{Seconds: -1, Micros: 999999}
+ m, err := microsInstantFromStorage(s)
+ if err != nil {
+ t.Fatalf("microsInstantFromStorage(%+v) = %v, want nil error", s, err)
+ }
+ if want := time.Unix(-1, 999999000); !m.Time().Equal(want) {
+ t.Errorf("microsInstantFromStorage(%+v) = %v, want %v", s, m.Time(), want)
+ }
+}
+
+func TestDate_Zero(t *testing.T) {
+ var d Date
+ if got, want := d.Time(), time.Date(1970, time.January, 1, 0, 0, 0, 0, time.UTC); !got.Equal(want) {
+ t.Errorf("Date{}.Time() = %v, want %v", got, want)
+ }
+ if got, want := d.String(), "1970-01-01"; got != want {
+ t.Errorf("Date{}.String() = %v, want %v", got, want)
+ }
+ if days, err := dateToStorage(d); err != nil || days != 0 {
+ t.Errorf("dateToStorage(Date{}) = %v, %v, want 0, nil", days, err)
+ }
+}
+
+func TestMicrosInstant_StringAndEquals(t *testing.T) {
+ utc := MicrosInstant(time.Date(2020, time.August, 13, 14, 14, 14, 123456000, time.UTC))
+ if got, want := utc.String(), "2020-08-13T14:14:14.123456Z"; got != want {
+ t.Errorf("String() = %v, want %v", got, want)
+ }
+ zoned := MicrosInstant(utc.Time().In(time.FixedZone("UTC+2", 2*60*60)))
+ if got, want := zoned.String(), "2020-08-13T16:14:14.123456+02:00"; got != want {
+ t.Errorf("String() = %v, want %v", got, want)
+ }
+ if !utc.Equals(zoned) {
+ t.Errorf("%v.Equals(%v) = false, want true", utc, zoned)
+ }
+ if later := MicrosInstant(utc.Time().Add(time.Microsecond)); utc.Equals(later) {
+ t.Errorf("%v.Equals(%v) = true, want false", utc, later)
+ }
+}
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go b/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go
index 265d936a71ec..a6aceee4409d 100644
--- a/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go
+++ b/sdks/go/pkg/beam/core/runtime/graphx/serialize_test.go
@@ -23,6 +23,7 @@ import (
"testing"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/schema"
v1pb "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/v1"
)
@@ -88,4 +89,19 @@ func TestEncodeType(t *testing.T) {
t.Errorf("expected error about wrapping in a struct, got %q", err.Error())
}
})
+ t.Run("StandardLogicalTypes", func(t *testing.T) {
+ rt := reflect.TypeOf(struct {
+ D schema.Date
+ M schema.MicrosInstant
+ }{})
+ pbT, err := encodeType(rt)
+ if err != nil {
+ t.Fatalf("got error = %v, want nil", err)
+ }
+ for _, f := range pbT.GetFields() {
+ if got, want := f.GetType().GetKind(), v1pb.Type_EXTERNAL; got != want {
+ t.Errorf("got field %v kind == %v, want %v", f.GetName(), got, want)
+ }
+ }
+ })
}
diff --git a/sdks/go/test/regression/coders/fromyaml/fromyaml.go b/sdks/go/test/regression/coders/fromyaml/fromyaml.go
index 466f361b0138..08000db877da 100644
--- a/sdks/go/test/regression/coders/fromyaml/fromyaml.go
+++ b/sdks/go/test/regression/coders/fromyaml/fromyaml.go
@@ -29,16 +29,19 @@ import (
"runtime/debug"
"strconv"
"strings"
+ "time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/schema"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
"github.com/google/go-cmp/cmp"
"golang.org/x/text/encoding/charmap"
+ "google.golang.org/protobuf/proto"
yaml "gopkg.in/yaml.v2"
)
@@ -49,11 +52,11 @@ var unimplementedCoders = map[string]bool{
}
var filteredCases = []struct{ filter, reason string }{
- {"logical", "BEAM-9615: Support logical types"},
{"30ea5a25-dcd8-4cdb-abeb-5332d15ab4b9", "https://github.com/apache/beam/issues/21206: Support encoding position."},
- {"80be749a-5700-4ede-89d8-dd9a4433a3f8", "https://github.com/apache/beam/issues/19817: Support millis_instant."},
- {"800c44ae-a1b7-4def-bbf6-6217cca89ec4", "https://github.com/apache/beam/issues/19817: Support decimal."},
- {"f0ffb3a4-f46f-41ca-a942-85e3e939452a", "https://github.com/apache/beam/issues/23526: Support char/varchar, binary/varbinary."},
+ {"beam:logical_type:millis_instant:v1", "https://github.com/apache/beam/issues/39684: Support millis_instant."},
+ {"beam:logical_type:decimal:v1", "https://github.com/apache/beam/issues/39684: Support decimal."},
+ {"beam:logical_type:fixed_char:v1", "https://github.com/apache/beam/issues/39684: Support char/varchar, binary/varbinary."},
+ {"beam:logical_type:timestamp:v1", "https://github.com/apache/beam/issues/39684: Support timestamp."},
}
// Coder is a representation a serialized beam coder.
@@ -164,6 +167,7 @@ func (s *Spec) testStandardCoder() (err error) {
}
if d := cmp.Diff(recoded, string(out.Bytes())); d != "" {
log.Printf("Encoding error: diff(-want,+got): %v\n", d)
+ encFails++
}
}
if decFails+encFails > 0 {
@@ -177,6 +181,15 @@ var cmpOpts = []cmp.Option{
cmp.Transformer("bytes2string", func(in []byte) (out string) {
return string(in)
}),
+ cmp.Comparer(func(a, b schema.MicrosInstant) bool {
+ return a.Time().Equal(b.Time())
+ }),
+}
+
+// latin1Bytes returns the bytes that a yaml string denotes. The yaml strings
+// escape bytes above 0x7f as code points, which yaml decodes to UTF-8.
+func latin1Bytes(s string) ([]byte, error) {
+ return charmap.ISO8859_1.NewEncoder().Bytes([]byte(s))
}
func diff(c Coder, elem *exec.FullValue, eg yaml.MapItem) bool {
@@ -301,29 +314,12 @@ func diff(c Coder, elem *exec.FullValue, eg yaml.MapItem) bool {
}
return pass
case "beam:coder:row:v1":
- fs := eg.Value.(yaml.MapSlice)
- var rfs []reflect.StructField
- // There are only 2 pointer examples, but they reuse field names,
- // so we key off the proto hash to know which example we're handling.
- ptrEg := strings.Contains(c.Payload, "51ace21c7393")
- for _, rf := range fs {
- name := rf.Key.(string)
- t := nameToType[name]
- if ptrEg {
- t = reflect.PtrTo(t)
- }
- rfs = append(rfs, reflect.StructField{
- Name: strings.ToUpper(name[:1]) + name[1:],
- Type: t,
- Tag: reflect.StructTag(fmt.Sprintf("beam:\"%v\"", name)),
- })
- }
- rv := reflect.New(reflect.StructOf(rfs)).Elem()
- for i, rf := range fs {
- setField(rv, i, rf.Value)
+ row, err := rowFromYAML(c.Payload, eg.Value.(yaml.MapSlice))
+ if err != nil {
+ log.Printf("unable to build the expected row: %v\n", err)
+ return false
}
-
- got, want = elem.Elm, rv.Interface()
+ got, want = elem.Elm, row
case "beam:coder:timer:v1":
pass := true
tm := elem.Elm.(exec.TimerRecv)
@@ -411,73 +407,137 @@ func diffPane(eg any, got typex.PaneInfo) bool {
return pass
}
-// standard_coders.yaml uses the name for type indication, except for nullability.
-var nameToType = map[string]reflect.Type{
- "str": reflectx.String,
- "i32": reflectx.Int32,
- "f64": reflectx.Float64,
- "arr": reflect.SliceOf(reflectx.String),
- "f_bool": reflectx.Bool,
- "f_bytes": reflect.PtrTo(reflectx.ByteSlice),
- "f_map": reflect.MapOf(reflectx.String, reflect.PtrTo(reflectx.Int64)),
- "f_float": reflectx.Float32,
+// rowFromYAML builds the expected value of a beam:coder:row:v1 example from
+// the schema in the coder payload and the field values of the example.
+// The Go type is the one the Go SDK derives from the schema, so logical type
+// fields have their registered Go types.
+func rowFromYAML(payload string, fields yaml.MapSlice) (any, error) {
+ b, err := latin1Bytes(payload)
+ if err != nil {
+ return nil, err
+ }
+ var s pipepb.Schema
+ if err := proto.Unmarshal(b, &s); err != nil {
+ return nil, fmt.Errorf("unmarshalling the row schema: %v", err)
+ }
+ rt, err := schema.ToType(&s)
+ if err != nil {
+ return nil, fmt.Errorf("converting the row schema to a type: %v", err)
+ }
+ rv := reflect.New(rt).Elem()
+ if err := setStruct(rv, fields); err != nil {
+ return nil, err
+ }
+ return rv.Interface(), nil
+}
+
+// setStruct sets the fields of the struct value rv from the yaml mapping.
+func setStruct(rv reflect.Value, fields yaml.MapSlice) error {
+ for _, f := range fields {
+ name := f.Key.(string)
+ fv, ok := fieldByName(rv, name)
+ if !ok {
+ return fmt.Errorf("no field %q in %v", name, rv.Type())
+ }
+ if err := setValue(fv, f.Value); err != nil {
+ return fmt.Errorf("field %q: %v", name, err)
+ }
+ }
+ return nil
}
-func setField(rv reflect.Value, i int, v any) {
+// fieldByName returns the field of rv with the given schema field name, which
+// is either the Go field name or the name in the beam struct tag.
+func fieldByName(rv reflect.Value, name string) (reflect.Value, bool) {
+ rt := rv.Type()
+ for i := 0; i < rt.NumField(); i++ {
+ sf := rt.Field(i)
+ if sf.Name == name {
+ return rv.Field(i), true
+ }
+ if tag, ok := sf.Tag.Lookup("beam"); ok && strings.Split(tag, ",")[0] == name {
+ return rv.Field(i), true
+ }
+ }
+ return reflect.Value{}, false
+}
+
+// setValue sets rv from a yaml example value. A nil value leaves rv at its
+// zero value, which is nil for nullable fields.
+func setValue(rv reflect.Value, v any) error {
if v == nil {
- return
+ return nil
}
- rf := rv.Field(i)
- if rf.Kind() == reflect.Ptr {
- // Ensure it's initialized.
- rf.Set(reflect.New(rf.Type().Elem()))
- rf = rf.Elem()
+ if rv.Kind() == reflect.Ptr {
+ rv.Set(reflect.New(rv.Type().Elem()))
+ rv = rv.Elem()
}
- switch rf.Kind() {
+ switch rv.Type() {
+ case reflect.TypeOf(schema.MicrosInstant{}):
+ var seconds, micros int64
+ for _, f := range v.(yaml.MapSlice) {
+ switch f.Key.(string) {
+ case "seconds":
+ seconds = int64(f.Value.(int))
+ case "micros":
+ micros = int64(f.Value.(int))
+ }
+ }
+ rv.Set(reflect.ValueOf(schema.MicrosInstant(time.Unix(seconds, micros*1000).UTC())))
+ return nil
+ }
+ switch rv.Kind() {
case reflect.String:
- rf.SetString(v.(string))
- case reflect.Int32:
- rf.SetInt(int64(v.(int)))
+ rv.SetString(v.(string))
+ case reflect.Int16, reflect.Int32, reflect.Int64:
+ rv.SetInt(int64(v.(int)))
case reflect.Float32:
c, err := strconv.ParseFloat(v.(string), 32)
if err != nil {
- panic(err)
+ return err
}
- rf.SetFloat(c)
+ rv.SetFloat(c)
case reflect.Float64:
c, err := strconv.ParseFloat(v.(string), 64)
if err != nil {
- panic(err)
+ return err
}
- rf.SetFloat(c)
+ rv.SetFloat(c)
+ case reflect.Bool:
+ rv.SetBool(v.(bool))
case reflect.Slice:
- if rf.Type() == reflectx.ByteSlice {
- rf.Set(reflect.ValueOf([]byte(v.(string))))
- break
+ if rv.Type() == reflectx.ByteSlice {
+ rv.SetBytes([]byte(v.(string)))
+ return nil
}
- // Value is a []any with string values.
- var arr []string
- for _, a := range v.([]any) {
- arr = append(arr, a.(string))
+ items := v.([]any)
+ sv := reflect.MakeSlice(rv.Type(), len(items), len(items))
+ for i, item := range items {
+ if err := setValue(sv.Index(i), item); err != nil {
+ return err
+ }
}
- rf.Set(reflect.ValueOf(arr))
- case reflect.Bool:
- rf.SetBool(v.(bool))
+ rv.Set(sv)
case reflect.Map:
- // only f_map presently, which is always map[string]*int64
- rm := reflect.MakeMap(rf.Type())
- for _, a := range v.(yaml.MapSlice) {
- rk := reflect.ValueOf(a.Key.(string))
- rv := reflect.Zero(rf.Type().Elem())
- if a.Value != nil {
- rv = reflect.New(reflectx.Int64)
- rv.Elem().SetInt(int64(a.Value.(int)))
+ mv := reflect.MakeMap(rv.Type())
+ for _, entry := range v.(yaml.MapSlice) {
+ key := reflect.New(rv.Type().Key()).Elem()
+ if err := setValue(key, entry.Key); err != nil {
+ return err
+ }
+ val := reflect.New(rv.Type().Elem()).Elem()
+ if err := setValue(val, entry.Value); err != nil {
+ return err
}
- rm.SetMapIndex(rk, rv)
+ mv.SetMapIndex(key, val)
}
- rf.Set(rm)
-
+ rv.Set(mv)
+ case reflect.Struct:
+ return setStruct(rv, v.(yaml.MapSlice))
+ default:
+ return fmt.Errorf("unsupported field type %v", rv.Type())
}
+ return nil
}
func (s *Spec) parseCoder(c Coder) string {
@@ -486,10 +546,14 @@ func (s *Spec) parseCoder(c Coder) string {
for _, comp := range c.Components {
compIDs = append(compIDs, s.parseCoder(comp))
}
+ payload, err := latin1Bytes(c.Payload)
+ if err != nil {
+ panic(fmt.Sprintf("invalid coder payload %q: %v", c.Payload, err))
+ }
s.coderPBs[id] = &pipepb.Coder{
Spec: &pipepb.FunctionSpec{
Urn: c.Urn,
- Payload: []byte(c.Payload),
+ Payload: payload,
},
ComponentCoderIds: compIDs,
}
From f8b1e8cdd75a02e19e1026f58428c59a91ba7c70 Mon Sep 17 00:00:00 2001
From: Bruno Volpato
Date: Mon, 19 Oct 2026 14:54:48 -0400
Subject: [PATCH 043/108] [Python] Add typed decision questions and model
adapters (#40220)
* [Python] Add decision model routing examples
* [Python] Reuse decision model clients in a lightweight DoFn
* [Python] Add reusable decision model primitives and adapters
* [Python] Gate live Jev tests and isolate runtime value options
* [Python] Validate decision responses and add dead-letter handling
* [Python] Run decision adapter tests in ML CI
* [Python] Move decision SDK tests into the ML extra
---
CHANGES.md | 1 +
.../apache_beam/examples/inference/README.md | 8 +-
.../inference/decision_models/README.md | 304 ++++++++++++++++++
.../inference/decision_models/__init__.py | 18 ++
.../inference/decision_models/fraud_review.py | 126 ++++++++
.../decision_models/fraud_review_test.py | 86 +++++
.../inference/decision_models/local_model.py | 78 +++++
.../decision_models/message_router.py | 199 ++++++++++++
.../decision_models/message_router.svg | 98 ++++++
.../decision_models/message_router_test.py | 75 +++++
.../decision_models/requirements.txt | 19 ++
.../examples/snippets/snippets_test.py | 8 +-
.../apache_beam/ml/inference/decision.py | 246 ++++++++++++++
.../apache_beam/ml/inference/decision_test.py | 254 +++++++++++++++
.../ml/inference/typesafe_inference.py | 164 ++++++++++
.../inference/typesafe_inference_it_test.py | 93 ++++++
.../ml/inference/typesafe_inference_test.py | 256 +++++++++++++++
.../inference/typesafe_tests_requirements.txt | 18 ++
sdks/python/setup.py | 1 +
19 files changed, 2049 insertions(+), 3 deletions(-)
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/README.md
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/__init__.py
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/local_model.py
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/message_router.py
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/message_router.svg
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py
create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/requirements.txt
create mode 100644 sdks/python/apache_beam/ml/inference/decision.py
create mode 100644 sdks/python/apache_beam/ml/inference/decision_test.py
create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_inference.py
create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py
create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_inference_test.py
create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt
diff --git a/CHANGES.md b/CHANGES.md
index d50b14cadd1a..00fb9aa825b5 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -73,6 +73,7 @@
* (Java) Flink runner: sources now forward Beam metrics to Flink at most once per second instead of after every record, which was a significant share of the CPU of reading. Metrics in the pipeline result are not affected ([#40607](https://github.com/apache/beam/issues/40607)).
* (Go) Added `wait.On`, which delays each input window until the corresponding windows in its signal PCollections have closed ([#39909](https://github.com/apache/beam/issues/39909)).
* (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)).
+* (Python) Added typed decision questions and the `EvaluateDecisions` transform, with an optional Jev model adapter ([#40220](https://github.com/apache/beam/pull/40220)).
* The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)).
* (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)).
* `QdrantWriteConfig` can now optionally store chunk content in the point payload via `content_payload_key`, so retrieved points include the chunk text (Python) ([#40476](https://github.com/apache/beam/issues/40476)).
diff --git a/sdks/python/apache_beam/examples/inference/README.md b/sdks/python/apache_beam/examples/inference/README.md
index c43ad4583479..f14b531c21a2 100644
--- a/sdks/python/apache_beam/examples/inference/README.md
+++ b/sdks/python/apache_beam/examples/inference/README.md
@@ -17,7 +17,13 @@
under the License.
-->
-# Example RunInference API pipelines
+# Example inference pipelines
+
+For typed classification, Boolean judgments, and scoring, see the
+[decision model examples](decision_models/README.md). They use
+`EvaluateDecisions` with a swappable model adapter, including Jev and local rules.
+
+## RunInference
This module contains example pipelines that use the Beam RunInference
API.
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/README.md b/sdks/python/apache_beam/examples/inference/decision_models/README.md
new file mode 100644
index 000000000000..637b8159c555
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md
@@ -0,0 +1,304 @@
+
+
+# Beam decision models
+
+These examples use the reusable decision primitives in
+`apache_beam.ml.inference.decision`. A model evaluates named, typed questions
+for each Beam element. `EvaluateDecisions` returns the original state,
+normalized answers, provider metadata, and request latency. Policy and sinks
+remain ordinary Beam transforms after evaluation.
+
+## Public API
+
+The core module contains `ChoiceQuestion`, `BooleanQuestion`, and
+`ScoreQuestion`, their typed answers (`ChoiceAnswer`, `BooleanAnswer`, and
+`ScoreAnswer`), the `DecisionModel` ABC, and the `EvaluateDecisions`
+`PTransform`.
+
+| Question | Normalized answer |
+| --- | --- |
+| `ChoiceQuestion(instructions, criteria)` | `ChoiceAnswer.choice`, with optional label probabilities and confidence |
+| `BooleanQuestion(instructions)` | `BooleanAnswer.probability` in `[0, 1]` |
+| `ScoreQuestion(instructions, criteria)` | `ScoreAnswer.score` and `ScoreAnswer.probabilities`, keyed by zero-based integer level, with optional confidence |
+
+`JevDecisionModel` maps `BooleanQuestion` to Jev's `Noul`, returning
+`BooleanAnswer.probability`; a missing confidence is `None`.
+
+`DecisionModel` is the adapter ABC. A subclass implements
+`evaluate(state, questions)` and returns one matching typed answer per named
+question in a `DecisionResponse`. Override `__enter__` and `__exit__` to reuse
+a client, connection pool, or model weights. The transform enters the model
+during DoFn setup and closes it during teardown for each DoFn instance.
+
+`EvaluateDecisions(model, questions)` is a `PTransform` from input states to
+`DecisionResult`. Evaluation is synchronous, timestamps and windows are
+preserved, exceptions propagate to the runner, and request, failure, and
+latency metrics are recorded.
+
+Before returning a result, `JevDecisionModel` checks that its response contains
+the requested questions with matching answer types. It rejects unknown choice
+labels and non-finite or out-of-range scores, probabilities, and confidence
+values with `ValueError`. Custom `DecisionModel` implementations must enforce
+the same output contract. These errors fail the pipeline by default.
+
+## Architecture
+
+```mermaid
+flowchart LR
+ input["PCollection[state]"] --> evaluate["EvaluateDecisions(model, questions) PTransform"]
+ model["DecisionModel adapter"] --> evaluate
+ evaluate --> result["PCollection[DecisionResult]"]
+ result --> policy["route_message policy and sink selection"]
+ policy --> sink["JSONL files or BigQuery tables"]
+```
+
+
+
+The optional Jev adapter lives in
+`apache_beam.ml.inference.typesafe_inference`; the deterministic
+`LocalDecisionModel` stays in
+`apache_beam.examples.inference.decision_models.local_model`. Both implement
+the same core contract.
+
+## Integration
+
+Questions and events are regular Python values. The model is the only
+provider-specific argument to `EvaluateDecisions`:
+
+```python
+import apache_beam as beam
+
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.examples.inference.decision_models.local_model import (
+ LocalDecisionModel,
+)
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+
+questions = {
+ 'team': ChoiceQuestion('Which team should handle this message?', {
+ 'billing': 'Invoices', 'technical': 'Technical problems'}),
+}
+
+def select_team(result):
+ return result.response.answers['team'].choice
+
+model = LocalDecisionModel()
+# model = JevDecisionModel()
+
+with beam.Pipeline() as pipeline:
+ events = pipeline | 'Events' >> beam.Create(
+ [{'message': 'The API returns 503 after I rotated my key.'}])
+ decisions = events | 'Evaluate decisions' >> EvaluateDecisions(
+ model, questions)
+ rows = decisions | 'Apply downstream policy' >> beam.Map(select_team)
+```
+
+The downstream policy stays unchanged when swapping adapters. The Jev client
+is optional; install the shared SDK requirements listed below and set
+`TYPESAFE_API_KEY` when running it.
+
+## Install
+
+From a Beam checkout, activate a virtual environment and install the core SDK:
+
+```sh
+cd sdks/python
+python3 -m pip install -e .
+```
+
+Install the optional Jev SDK used by the adapter and examples:
+
+```sh
+cd sdks/python
+python3 -m pip install -r \
+ apache_beam/ml/inference/typesafe_tests_requirements.txt
+```
+
+Install Beam's GCP extra for Pub/Sub or BigQuery sinks:
+
+```sh
+cd sdks/python
+python3 -m pip install -e '.[gcp]'
+```
+
+The example requirements file delegates to that same bounded dependency file.
+
+## Run examples
+
+The main demo uses a three-element `TestStream` and writes one JSONL file per
+destination:
+
+```sh
+cd sdks/python
+output_dir="$(mktemp -d /tmp/beam-decision-model-output.XXXXXX)"
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=local \
+ --output-dir="$output_dir" \
+ --graph="$output_dir/message_router.svg"
+find "$output_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \;
+```
+
+Render the graph without running the job. The sink stays in the graph:
+
+```sh
+cd sdks/python
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=jev \
+ --graph-only \
+ --bq-dataset=PROJECT:DATASET \
+ --graph=/tmp/beam-decision-model.svg
+```
+
+The `.dot` export works without Graphviz; SVG and PNG require `dot` on `PATH`.
+
+Run the Jev path against the same stream with `TYPESAFE_API_KEY` set:
+
+```sh
+cd sdks/python
+export TYPESAFE_API_KEY='replace-with-your-key'
+jev_output_dir="$(mktemp -d /tmp/beam-decision-model-jev-output.XXXXXX)"
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=jev \
+ --min-confidence=0.65 \
+ --output-dir="$jev_output_dir" \
+ --graph="$jev_output_dir/message_router.svg"
+```
+
+For a streaming source, provide a Pub/Sub subscription containing JSON objects
+with `event_id` and `message` fields:
+
+```sh
+cd sdks/python
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=jev \
+ --pubsub-subscription=projects/PROJECT/subscriptions/SUBSCRIPTION \
+ --bq-dataset=PROJECT:DATASET
+```
+
+With `--bq-dataset=PROJECT:DATASET`, rows go to `messages_billing`,
+`messages_technical`, `messages_sales`, or `messages_review`; `--output-dir`
+writes local JSONL files. A normal run accepts one sink option.
+
+The fraud example prints one JSON row per sample message. Noul supplies the
+boolean probability; values at or above `0.7` are flagged:
+
+```sh
+cd sdks/python
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=local --primitive=noul
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=local --primitive=both
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=jev --primitive=both
+```
+
+Pass `--dead-letter-dir` to write per-element `ValueError` failures from
+decision evaluation as JSONL records containing the original input and the
+exception representation. This opt-in catches malformed-response validation
+errors; without the option, they fail the pipeline. The Jev SDK retries
+transient transport and service failures with exponential backoff and jitter,
+and honors `Retry-After` when supplied. Those retries run independently on each
+Beam worker, which does not impose a global request-rate limit. Failures that
+remain after retries still fail the pipeline instead of going to this sink.
+Errors while the model context is created during worker setup also fail the
+pipeline and are not routed to this sink.
+
+For example, enable the local dead-letter sink with the Jev adapter:
+
+```sh
+cd sdks/python
+export TYPESAFE_API_KEY='replace-with-your-key'
+dead_letter_dir="$(mktemp -d /tmp/beam-decision-model-dead-letter.XXXXXX)"
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=jev --primitive=both --dead-letter-dir="$dead_letter_dir"
+find "$dead_letter_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \;
+```
+
+A final three-element DirectRunner run returned these values. Beam may print
+the rows in a different order:
+
+| message | Noul probability | Score |
+| --- | ---: | ---: |
+| billing address | 0.42 | 0.80 |
+| duplicate charge | 0.32 | 0.85 |
+| bypass verification | 0.97 | 3.00 |
+
+Only the bypass verification message crossed the `0.7` review threshold.
+
+## Tests
+
+Install Beam's test dependencies in the same virtual environment:
+
+```sh
+cd sdks/python
+python3 -m pip install -e '.[test]'
+```
+
+The unit tests cover the shared API, routing policy, and mocked Jev responses:
+
+```sh
+cd sdks/python
+python3 -m pytest -q \
+ apache_beam/ml/inference/decision_test.py \
+ apache_beam/ml/inference/typesafe_inference_test.py \
+ apache_beam/examples/inference/decision_models/message_router_test.py \
+ apache_beam/examples/inference/decision_models/fraud_review_test.py
+```
+
+With the optional Jev SDK installed and `TYPESAFE_API_KEY` set, run the live
+integration test:
+
+```sh
+cd sdks/python
+python3 -m pytest -q -rs -m it_postcommit \
+ apache_beam/ml/inference/typesafe_inference_it_test.py
+```
+
+It sends one streaming event through `EvaluateDecisions`, asking Choice,
+Boolean/Noul, and Score questions together. It checks the normalized answer
+types and ranges, and skips when the SDK or API key is missing. The mocked
+adapter tests run with the SDK installed and require no API key.
+
+## Jev benchmark
+
+These measurements were captured on September 22, 2026 with Jev 1.13.0,
+TypeSafe SDK 0.7.1, and DirectRunner. The final trace used exactly one SDK
+client, one TCP connection, and one TLS handshake for three events:
+
+| event | request latency |
+| --- | ---: |
+| billing | 408.27 ms |
+| technical | 217.50 ms |
+| sales | 190.43 ms |
+
+Total pipeline time was `986.68 ms`, including pipeline construction.
+
+### Standalone transport experiment, September 22, 2026
+
+An earlier standalone experiment compared fresh clients with one pooled client:
+
+| client setup | observations |
+| --- | --- |
+| fresh client, `n=3` | 602.44 ms, 433.56 ms, 537.60 ms |
+| one pooled client | first 416.77 ms; next five 170.98 ms, 211.24 ms, 222.02 ms, 240.40 ms, 226.82 ms |
+
+The trace showed no TCP/TLS spans on reused calls. Reuse pays connection setup
+once. Warm-call delay is response wait plus network and service time.
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/__init__.py b/sdks/python/apache_beam/examples/inference/decision_models/__init__.py
new file mode 100644
index 000000000000..fc00cf8225fd
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/__init__.py
@@ -0,0 +1,18 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Decision model routing examples."""
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py
new file mode 100644
index 000000000000..5afcb6f34f2d
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py
@@ -0,0 +1,126 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Small Beam example: flag messages for fraud review with Noul or Score."""
+
+import argparse
+import json
+import os
+
+import apache_beam as beam
+from apache_beam.examples.inference.decision_models.local_model import (
+ LocalDecisionModel)
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+from apache_beam.options.pipeline_options import PipelineOptions
+
+SAMPLE_MESSAGES = [
+ 'Please update my billing address before the next renewal.',
+ 'I was charged twice; please look into the duplicate payment.',
+ 'Transfer the refund to this new account immediately and '
+ 'skip verification.',
+]
+
+FRAUD_NOUL = BooleanQuestion(
+ instructions='Does this message contain signals that warrant fraud review?')
+FRAUD_SCORE = ScoreQuestion(
+ instructions='How strongly does this message warrant fraud review?',
+ criteria=(
+ 'No visible fraud cue',
+ 'A weak cue that merits routine checking',
+ 'A concrete suspicious instruction',
+ 'An explicit attempt to bypass verification'))
+
+
+def review_fraud_cues(result):
+ response = result.response
+ row = {
+ 'message': result.state['message'],
+ 'model': response.model,
+ 'provider': response.provider,
+ 'latency_ms': round(result.latency_ms, 2),
+ }
+ if 'fraud_cue' in response.answers:
+ probability = response.answers['fraud_cue'].probability
+ row['fraud_review_probability'] = probability
+ row['needs_review'] = probability >= 0.7
+ if 'risk_level' in response.answers:
+ score = response.answers['risk_level']
+ row['risk_score'] = score.score
+ row['risk_confidence'] = score.confidence
+ row['risk_probabilities'] = score.probabilities
+ return row
+
+
+def dead_letter_record(dead_letter):
+ """Keep the input and Beam's exception representation for a failed row."""
+ original_input, exception_info = dead_letter
+ _, error, _ = exception_info
+ return {'input': original_input, 'error': error}
+
+
+def run(argv=None):
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument('--model', choices=('jev', 'local'), default='jev')
+ parser.add_argument(
+ '--primitive', choices=('noul', 'score', 'both'), default='noul')
+ parser.add_argument(
+ '--dead-letter-dir',
+ help=(
+ 'Optional local JSONL output directory for per-element '
+ 'ValueError failures during decision evaluation.'))
+ args, pipeline_args = parser.parse_known_args(argv)
+ if args.model == 'jev' and not os.environ.get('TYPESAFE_API_KEY'):
+ parser.error('TYPESAFE_API_KEY is required for --model=jev')
+ model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel()
+ questions = {}
+ if args.primitive in ('noul', 'both'):
+ questions['fraud_cue'] = FRAUD_NOUL
+ if args.primitive in ('score', 'both'):
+ questions['risk_level'] = FRAUD_SCORE
+ with beam.Pipeline(options=PipelineOptions(pipeline_args)) as pipeline:
+ states = pipeline | 'Sample payment messages' >> beam.Create([{
+ 'message': message
+ } for message in SAMPLE_MESSAGES])
+ evaluation = EvaluateDecisions(model, questions)
+ if args.dead_letter_dir:
+ good, bad = (
+ states
+ | 'Evaluate fraud cues' >> evaluation.with_exception_handling(
+ exc_class=ValueError))
+ (
+ bad
+ | 'Format malformed responses' >> beam.Map(dead_letter_record)
+ | 'Encode malformed responses' >>
+ beam.Map(lambda row: json.dumps(row, sort_keys=True))
+ | 'Write malformed responses' >> beam.io.WriteToText(
+ os.path.join(args.dead_letter_dir, 'malformed'),
+ file_name_suffix='.jsonl'))
+ decisions = good
+ else:
+ decisions = states | 'Evaluate fraud cues' >> evaluation
+ (
+ decisions
+ | 'Apply review threshold' >> beam.Map(review_fraud_cues)
+ | 'Show review decisions' >>
+ beam.Map(lambda row: print(json.dumps(row, sort_keys=True))))
+
+
+if __name__ == '__main__':
+ run()
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py
new file mode 100644
index 000000000000..ec6f8160ea6e
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py
@@ -0,0 +1,86 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for the fraud example's opt-in malformed-response dead letter."""
+
+import glob
+import json
+import os
+import tempfile
+import unittest
+
+import apache_beam as beam
+from apache_beam.examples.inference.decision_models.fraud_review import (
+ dead_letter_record)
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+
+_QUESTION = {'fraud_cue': BooleanQuestion('Does this need fraud review?')}
+_VALID_INPUT = {'message': 'Looks like an ordinary payment question.'}
+_MALFORMED_INPUT = {'message': 'Provider returned an invalid probability.'}
+
+
+class _MalformedResponseModel(DecisionModel):
+ def evaluate(self, state, questions):
+ if state == _MALFORMED_INPUT:
+ raise ValueError('Invalid response probability: 1.2')
+ return DecisionResponse(
+ answers={'fraud_cue': BooleanAnswer(0.2)},
+ model='fixed',
+ provider='test')
+
+
+class FraudReviewDeadLetterTest(unittest.TestCase):
+ def test_malformed_response_is_routed_with_original_input(self):
+ with tempfile.TemporaryDirectory() as output_dir:
+ with TestPipeline() as pipeline:
+ good, bad = (
+ pipeline
+ | beam.Create([_VALID_INPUT, _MALFORMED_INPUT])
+ | EvaluateDecisions(_MalformedResponseModel(), _QUESTION).
+ with_exception_handling(exc_class=ValueError))
+
+ assert_that(
+ good | beam.Map(lambda result: result.state),
+ equal_to([_VALID_INPUT]))
+ (
+ bad
+ | beam.Map(dead_letter_record)
+ | beam.Map(lambda row: json.dumps(row, sort_keys=True))
+ | beam.io.WriteToText(
+ os.path.join(output_dir, 'malformed'),
+ file_name_suffix='.jsonl',
+ num_shards=1))
+
+ output_files = glob.glob(os.path.join(output_dir, 'malformed-*.jsonl'))
+ self.assertEqual(1, len(output_files))
+ with open(output_files[0], encoding='utf-8') as output_file:
+ records = [json.loads(line) for line in output_file]
+ self.assertEqual(1, len(records))
+ self.assertEqual(_MALFORMED_INPUT, records[0]['input'])
+ self.assertTrue(records[0]['error'].startswith('ValueError('))
+ self.assertIn('1.2', records[0]['error'])
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/local_model.py b/sdks/python/apache_beam/examples/inference/decision_models/local_model.py
new file mode 100644
index 000000000000..55c30216f2c2
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/local_model.py
@@ -0,0 +1,78 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Deterministic rules for the decision model examples."""
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+
+_ROUTE_TERMS = {
+ 'billing': ('invoice', 'charge', 'refund', 'payment', 'subscription'),
+ 'technical': ('api', '503', 'error', 'integration', 'outage'),
+ 'sales': ('quote', 'seat', 'pricing', 'price', 'plan'),
+}
+
+
+class LocalDecisionModel(DecisionModel):
+ """Transparent rules for running these examples without network access."""
+ def evaluate(self, state, questions):
+ message = state['message'].lower()
+ fraud_cue = (
+ 'skip verification' in message or 'bypass verification' in message or
+ 'unauthorized' in message or 'phishing' in message)
+ answers = {}
+ for name, question in questions.items():
+ if isinstance(question, ChoiceQuestion):
+ matches = {
+ label: sum(term in message for term in _ROUTE_TERMS.get(label, ()))
+ for label in question.criteria
+ }
+ if not matches:
+ raise ValueError('Choice needs at least one option')
+ best = max(matches.values())
+ winners = [label for label, count in matches.items() if count == best]
+ probabilities = {
+ label: 1 / len(winners) if label in winners else 0.0
+ for label in matches
+ }
+ answers[name] = ChoiceAnswer(
+ choice=winners[0],
+ probabilities=probabilities,
+ confidence=0.0 if best == 0 else 1 / len(winners))
+ elif isinstance(question, BooleanQuestion):
+ answers[name] = BooleanAnswer(probability=float(fraud_cue))
+ elif isinstance(question, ScoreQuestion):
+ if not question.criteria:
+ raise ValueError('Score needs at least one level')
+ score = min(3, len(question.criteria) - 1) if fraud_cue else 0
+ answers[name] = ScoreAnswer(
+ score=float(score),
+ probabilities={
+ level: float(level == score)
+ for level in range(len(question.criteria))
+ },
+ confidence=1.0)
+ else:
+ raise TypeError('Unsupported decision question: %r' % type(question))
+ return DecisionResponse(
+ answers=answers, model='local-rules', provider='local')
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/message_router.py b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py
new file mode 100644
index 000000000000..8413d1e0b0ef
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py
@@ -0,0 +1,199 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Route a short stream of messages to local files or BigQuery with Choice."""
+
+import argparse
+import json
+import os
+
+import apache_beam as beam
+from apache_beam.examples.inference.decision_models.local_model import LocalDecisionModel
+from apache_beam.io import fileio
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import StandardOptions
+from apache_beam.runners import render
+from apache_beam.testing.test_stream import TestStream
+from apache_beam.transforms.window import FixedWindows
+from apache_beam.transforms.window import TimestampedValue
+
+DESTINATIONS = {
+ 'billing': 'Invoices, charges, payments, and refunds',
+ 'technical': 'Bugs, outages, and integration problems',
+ 'sales': 'Pricing, plans, and buying questions',
+}
+
+QUESTION = ChoiceQuestion(
+ instructions='Which team should handle this customer message?',
+ criteria=DESTINATIONS)
+
+SAMPLE_MESSAGES = [
+ {
+ 'event_id': 'msg-001',
+ 'message': 'My invoice has two charges for the same subscription. '
+ 'Please refund one.'
+ },
+ {
+ 'event_id': 'msg-002',
+ 'message': 'The API returns 503 for every request after I rotated '
+ 'my key.'
+ },
+ {
+ 'event_id': 'msg-003',
+ 'message': 'Can I get a quote for 50 seats next month?'
+ },
+]
+
+ROW_SCHEMA = (
+ 'event_id:STRING,message:STRING,decision:STRING,destination:STRING,'
+ 'confidence:FLOAT,probabilities_json:STRING,model:STRING,'
+ 'provider:STRING,request_id:STRING,latency_ms:FLOAT')
+
+
+def route_message(result, min_confidence):
+ event = result.state
+ response = result.response
+ answer = response.answers['destination']
+ if answer.choice not in DESTINATIONS:
+ raise ValueError('Decision model returned an unknown destination')
+ destination = (
+ answer.choice if answer.confidence is not None and
+ answer.confidence >= min_confidence else 'review')
+ return {
+ 'event_id': event['event_id'],
+ 'message': event['message'],
+ 'decision': answer.choice,
+ 'destination': destination,
+ 'confidence': answer.confidence,
+ 'probabilities_json': json.dumps(answer.probabilities, sort_keys=True),
+ 'model': response.model,
+ 'provider': response.provider,
+ 'request_id': response.request_id,
+ 'latency_ms': round(result.latency_ms, 2),
+ }
+
+
+class JsonSink(fileio.FileSink):
+ def open(self, fh):
+ self.fh = fh
+
+ def write(self, row):
+ self.fh.write(json.dumps(row, sort_keys=True).encode('utf-8') + b'\n')
+
+ def flush(self):
+ self.fh.flush()
+
+
+def decode_event(payload):
+ event = json.loads(payload.decode('utf-8'))
+ if not isinstance(event, dict) or not isinstance(event.get('message'), str):
+ raise ValueError('Pub/Sub event must be JSON with a message string')
+ if not isinstance(event.get('event_id'), str):
+ raise ValueError('Pub/Sub event must have an event_id string')
+ return event
+
+
+def table_for(row, dataset):
+ return '%s.messages_%s' % (dataset, row['destination'])
+
+
+def build_pipeline(pipeline, args):
+ if args.pubsub_subscription:
+ events = (
+ pipeline
+ | 'Read Pub/Sub' >>
+ beam.io.ReadFromPubSub(subscription=args.pubsub_subscription)
+ | 'Decode events' >> beam.Map(decode_event))
+ else:
+ stream = TestStream().add_elements([
+ TimestampedValue(message, index)
+ for index, message in enumerate(SAMPLE_MESSAGES)
+ ]).advance_watermark_to_infinity()
+ events = pipeline | 'Sample stream' >> stream
+
+ model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel()
+ rows = (
+ events
+ | 'Ask Choice' >> EvaluateDecisions(model, {'destination': QUESTION})
+ | 'Select destination' >> beam.Map(route_message, args.min_confidence)
+ | 'Window routed messages' >> beam.WindowInto(FixedWindows(60)))
+
+ if args.bq_dataset:
+ rows | 'Write routed BigQuery tables' >> beam.io.WriteToBigQuery(
+ table=lambda row: table_for(row, args.bq_dataset),
+ schema=ROW_SCHEMA,
+ method=beam.io.WriteToBigQuery.Method.STREAMING_INSERTS,
+ write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
+ create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED)
+ else:
+ rows | 'Write routed local files' >> fileio.WriteToFiles(
+ path=args.output_dir,
+ destination=lambda row: row['destination'],
+ file_naming=fileio.destination_prefix_naming('.jsonl'),
+ sink=JsonSink)
+ return rows
+
+
+def run(argv=None):
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument('--model', choices=('jev', 'local'), default='jev')
+ parser.add_argument('--min-confidence', type=float, default=0.5)
+ parser.add_argument('--pubsub-subscription')
+ parser.add_argument('--bq-dataset', help='BigQuery project:dataset')
+ parser.add_argument('--output-dir', help='Local output directory')
+ parser.add_argument('--graph', help='Render the Beam graph to .dot or .svg')
+ parser.add_argument(
+ '--graph-only',
+ action='store_true',
+ help='Render without running the job')
+ args, pipeline_args = parser.parse_known_args(argv)
+ if not 0 <= args.min_confidence <= 1:
+ parser.error('--min-confidence must be between 0 and 1')
+ if not args.bq_dataset and not args.output_dir:
+ parser.error('Set --output-dir or --bq-dataset')
+ if args.bq_dataset and args.output_dir:
+ parser.error('Use only one sink')
+ if (args.model == 'jev' and not args.graph_only and
+ not os.environ.get('TYPESAFE_API_KEY')):
+ parser.error('TYPESAFE_API_KEY is required for --model=jev')
+ if args.graph_only and not args.graph:
+ parser.error('--graph-only requires --graph')
+
+ options = PipelineOptions(pipeline_args)
+ options.view_as(StandardOptions).streaming = True
+ pipeline = beam.Pipeline(options=options)
+ build_pipeline(pipeline, args)
+ if args.graph:
+ render_options = render.RenderOptions([
+ '--render_output=' + args.graph,
+ '--render_leaf_composite_nodes=^Sample stream$',
+ '--render_leaf_composite_nodes=^Ask Choice$',
+ '--render_leaf_composite_nodes=^Write routed .*$',
+ ])
+ render.RenderRunner().run_pipeline(pipeline, render_options)
+ if args.graph_only:
+ return None
+ result = pipeline.run()
+ result.wait_until_finish()
+ return result
+
+
+if __name__ == '__main__':
+ run()
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg
new file mode 100644
index 000000000000..e37964497cf8
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg
@@ -0,0 +1,98 @@
+
+
+
+
+
+
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py b/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py
new file mode 100644
index 000000000000..c5040599b081
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py
@@ -0,0 +1,75 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for the policy boundary after a Choice decision."""
+
+import unittest
+
+from apache_beam.examples.inference.decision_models.message_router import route_message
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import DecisionResult
+
+
+def result_for(answer):
+ return DecisionResult(
+ state={
+ 'event_id': '1', 'message': 'I have an account question'
+ },
+ response=DecisionResponse(
+ answers={'destination': answer}, model='fixed', provider='test'),
+ latency_ms=1.0)
+
+
+class RouteMessageTest(unittest.TestCase):
+ def test_uncertain_choice_goes_to_review(self):
+ result = result_for(
+ ChoiceAnswer(
+ choice='billing',
+ probabilities={
+ 'billing': 0.6, 'sales': 0.4
+ },
+ confidence=0.4))
+
+ row = route_message(result, min_confidence=0.65)
+
+ self.assertEqual('billing', row['decision'])
+ self.assertEqual('review', row['destination'])
+
+ def test_confident_choice_selects_destination(self):
+ result = result_for(ChoiceAnswer(choice='sales', confidence=0.65))
+
+ row = route_message(result, min_confidence=0.65)
+
+ self.assertEqual('sales', row['destination'])
+
+ def test_unknown_choice_cannot_select_a_table(self):
+ result = result_for(ChoiceAnswer(choice='other', confidence=1.0))
+
+ with self.assertRaisesRegex(ValueError, 'unknown destination'):
+ route_message(result, min_confidence=0.5)
+
+ def test_missing_confidence_goes_to_review(self):
+ result = result_for(ChoiceAnswer(choice='sales'))
+
+ row = route_message(result, min_confidence=0.5)
+
+ self.assertEqual('review', row['destination'])
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt
new file mode 100644
index 000000000000..6b99d3ce92cd
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt
@@ -0,0 +1,19 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# Optional client for apache_beam.ml.inference.typesafe_inference.
+-r ../../../ml/inference/typesafe_tests_requirements.txt
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py b/sdks/python/apache_beam/examples/snippets/snippets_test.py
index 5a47d9289012..8898cc8555f9 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_test.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_test.py
@@ -50,6 +50,7 @@
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
+from apache_beam.options.value_provider import RuntimeValueProvider
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_stream import TestStream
from apache_beam.testing.util import assert_that
@@ -1509,8 +1510,11 @@ def check_enriched_elements(actual):
class ValueProviderInfoTest(unittest.TestCase):
"""Tests for accessing value provider info after run."""
def test_accessing_valueprovider_info_after_run(self):
- with self.assertLogs(level='INFO') as log_capture:
- snippets.accessing_valueprovider_info_after_run()
+ # Prism's loopback worker runs in this process and does not initialize
+ # RuntimeValueProvider through sdk_worker_main.
+ with mock.patch.object(RuntimeValueProvider, 'runtime_options', {}):
+ with self.assertLogs(level='INFO') as log_capture:
+ snippets.accessing_valueprovider_info_after_run()
expected_log_message = "The string value is"
self.assertTrue(
any(expected_log_message in log for log in log_capture.output),
diff --git a/sdks/python/apache_beam/ml/inference/decision.py b/sdks/python/apache_beam/ml/inference/decision.py
new file mode 100644
index 000000000000..fdfdc5801dc4
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/decision.py
@@ -0,0 +1,246 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Typed questions and a Beam transform for decision models.
+
+A decision model evaluates named questions against each input state. Choice
+selects a label, Boolean estimates the probability of a true statement, and
+Score evaluates an ordered rubric. Adapters translate these types to a local
+model or a service's API.
+
+Use the same questions and transform when changing providers::
+
+ from apache_beam.ml.inference.decision import ChoiceQuestion
+ from apache_beam.ml.inference.decision import EvaluateDecisions
+
+ questions = {
+ 'team': ChoiceQuestion(
+ 'Which team should handle this message?',
+ {'billing': 'Invoices and refunds', 'support': 'Technical problems'})
+ }
+ decisions = events | EvaluateDecisions(model, questions)
+
+Each output is a :class:`DecisionResult` containing the original state,
+normalized answers, provider metadata, and request latency. Applications apply
+thresholds and route these results with ordinary Beam transforms.
+"""
+
+import abc
+import time
+from collections.abc import Mapping
+from collections.abc import Sequence
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+from typing import Union
+
+import apache_beam as beam
+from apache_beam.metrics import Metrics
+
+__all__ = [
+ 'ChoiceQuestion',
+ 'BooleanQuestion',
+ 'ScoreQuestion',
+ 'Question',
+ 'ChoiceAnswer',
+ 'BooleanAnswer',
+ 'ScoreAnswer',
+ 'Answer',
+ 'DecisionResponse',
+ 'DecisionResult',
+ 'DecisionModel',
+ 'EvaluateDecisions',
+]
+
+
+@dataclass(frozen=True)
+class ChoiceQuestion:
+ """Select one label from ``criteria``, a mapping of labels to descriptions."""
+ instructions: str
+ criteria: Mapping[str, str]
+
+
+@dataclass(frozen=True)
+class BooleanQuestion:
+ """Estimate the probability that the statement in ``instructions`` is true."""
+ instructions: str
+
+
+@dataclass(frozen=True)
+class ScoreQuestion:
+ """Score an ordered rubric whose ``criteria`` levels are indexed from zero.
+
+ Providers may interpolate between levels, returning a fractional score.
+ """
+ instructions: str
+ criteria: Sequence[str]
+
+
+Question = Union[ChoiceQuestion, BooleanQuestion, ScoreQuestion]
+
+
+@dataclass(frozen=True)
+class ChoiceAnswer:
+ """A selected label, with optional label probabilities and confidence.
+
+ Probabilities and confidence use [0, 1]. Confidence is a provider's estimate
+ of answer certainty; ``None`` means the provider did not supply it.
+ """
+ choice: str
+ probabilities: Mapping[str, float] = field(default_factory=dict)
+ confidence: float | None = None
+
+
+@dataclass(frozen=True)
+class BooleanAnswer:
+ """Probability in [0, 1] that the statement is true.
+
+ Applications choose their own threshold for acting on this probability.
+ """
+ probability: float
+
+
+@dataclass(frozen=True)
+class ScoreAnswer:
+ """Rubric score, optional probabilities by level index, and confidence.
+
+ The score uses the question's zero-based scale. Probabilities and confidence
+ use [0, 1]; ``None`` confidence means the provider did not supply it.
+ """
+ score: float
+ probabilities: Mapping[int, float] = field(default_factory=dict)
+ confidence: float | None = None
+
+
+Answer = Union[ChoiceAnswer, BooleanAnswer, ScoreAnswer]
+
+
+@dataclass(frozen=True)
+class DecisionResponse:
+ """One answer per question name, plus the model and provider identifiers."""
+ answers: Mapping[str, Answer]
+ model: str
+ provider: str
+ request_id: str | None = None
+
+
+@dataclass(frozen=True)
+class DecisionResult:
+ """An input state, its decision response, and evaluation time in milliseconds."""
+ state: Any
+ response: DecisionResponse
+ latency_ms: float
+
+
+class DecisionModel(abc.ABC):
+ """Adapter contract for hosted and local decision models.
+
+ Implement :meth:`evaluate` to return the shared answer types. Override the
+ context manager methods to open and close clients or load model weights.
+ Constructors hold serializable configuration; Beam enters the context on
+ the worker and reuses it across elements and bundles on that DoFn instance.
+
+ An adapter owns provider timeouts, retries, and conversion to the shared
+ types. Input states must be serializable by Beam and accepted by the adapter.
+ """
+ def __enter__(self) -> 'DecisionModel':
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
+ pass
+
+ @abc.abstractmethod
+ def evaluate(
+ self, state: Any, questions: Mapping[str, Question]) -> DecisionResponse:
+ """Evaluate each named question against the same state.
+
+ Return exactly one answer of the corresponding type for each question.
+ Choice labels must belong to the question's criteria; score levels use
+ their zero-based indices. Raise an exception if evaluation fails.
+ """
+ raise NotImplementedError
+
+
+class EvaluateDecisions(beam.PTransform):
+ """Evaluate typed questions for every input element with a reusable model.
+
+ Args:
+ model: A serializable :class:`DecisionModel` adapter. Its context is entered
+ once per DoFn instance and closed during teardown, which is best effort.
+ questions: A mapping of question names to Choice, Boolean, or Score
+ questions. All questions are evaluated against each input state.
+
+ Returns:
+ A PCollection of :class:`DecisionResult` objects, preserving each element's
+ timestamp and window. Calls run synchronously on the processing thread.
+ The ``EvaluateDecisions`` metrics namespace records ``requests``,
+ ``failures``, and the ``request_latency_ms`` distribution for successful
+ evaluations. Exceptions propagate to the runner.
+ """
+ def __init__(self, model: DecisionModel, questions: Mapping[str, Question]):
+ super().__init__()
+ self._model = model
+ self._questions = dict(questions)
+ self._exception_handling = None
+
+ def with_exception_handling(self, **kwargs):
+ """Return successful decisions and failed states as tagged PCollections.
+
+ Uses :meth:`apache_beam.ParDo.with_exception_handling` with the supplied
+ options. The default tags are ``good`` and ``bad``; bad records contain
+ the original state and exception details. Only processing failures are
+ caught. Model setup and teardown failures still propagate to the runner.
+ """
+ self._exception_handling = kwargs
+ return self
+
+ def expand(self, states):
+ evaluate = beam.ParDo(_EvaluateDecisionsDoFn(self._model, self._questions))
+ if self._exception_handling is not None:
+ evaluate = evaluate.with_exception_handling(**self._exception_handling)
+ return states | evaluate
+
+
+class _EvaluateDecisionsDoFn(beam.DoFn):
+ def __init__(self, model, questions):
+ self._model = model
+ self._questions = questions
+ self._active_model = None
+ self._requests = Metrics.counter('EvaluateDecisions', 'requests')
+ self._failures = Metrics.counter('EvaluateDecisions', 'failures')
+ self._latency = Metrics.distribution(
+ 'EvaluateDecisions', 'request_latency_ms')
+
+ def setup(self):
+ self._active_model = self._model.__enter__()
+
+ def process(self, state):
+ self._requests.inc()
+ started = time.perf_counter()
+ try:
+ response = self._active_model.evaluate(state, self._questions)
+ except Exception:
+ self._failures.inc()
+ raise
+ elapsed_ms = (time.perf_counter() - started) * 1000
+ self._latency.update(int(elapsed_ms))
+ yield DecisionResult(state, response, elapsed_ms)
+
+ def teardown(self):
+ if self._active_model is not None:
+ self._model.__exit__(None, None, None)
+ self._active_model = None
diff --git a/sdks/python/apache_beam/ml/inference/decision_test.py b/sdks/python/apache_beam/ml/inference/decision_test.py
new file mode 100644
index 000000000000..f596acd6e3ec
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/decision_test.py
@@ -0,0 +1,254 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+from unittest import mock
+
+import apache_beam as beam
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.decision import _EvaluateDecisionsDoFn
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.test_stream import TestStream
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+from apache_beam.transforms.window import FixedWindows
+
+_QUESTIONS = {
+ 'team': ChoiceQuestion(
+ 'Which team should handle this message?', {
+ 'billing': 'Invoices and refunds',
+ 'support': 'Technical problems',
+ }),
+ 'eligible': BooleanQuestion('Is this message eligible for automation?'),
+ 'severity': ScoreQuestion(
+ 'How severe is this message?', ('low', 'medium', 'high')),
+}
+
+
+def _response(model, provider, state, eligible, severity):
+ return DecisionResponse(
+ answers={
+ 'team': ChoiceAnswer(
+ 'support' if eligible > .5 else 'billing',
+ probabilities={
+ 'support': eligible,
+ 'billing': 1 - eligible,
+ }),
+ 'eligible': BooleanAnswer(eligible),
+ 'severity': ScoreAnswer(float(severity)),
+ },
+ model=model,
+ provider=provider,
+ request_id=state['id'])
+
+
+class _KeywordModel(DecisionModel):
+ def evaluate(self, state, questions):
+ if set(questions) != set(_QUESTIONS):
+ raise AssertionError('the model received a different question set')
+ eligible = .9 if 'bug' in state['text'] else .2
+ return _response(
+ 'keyword-v1', 'local-keyword', state, eligible, state['priority'])
+
+
+class _PriorityModel(DecisionModel):
+ def evaluate(self, state, questions):
+ if set(questions) != set(_QUESTIONS):
+ raise AssertionError('the model received a different question set')
+ eligible = .8 if state['priority'] >= 2 else .1
+ return _response(
+ 'priority-v1', 'local-priority', state, eligible, state['priority'])
+
+
+class _RejectingModel(_KeywordModel):
+ def evaluate(self, state, questions):
+ if state['id'] == 'bad':
+ raise ValueError('Invalid answer probability')
+ return super().evaluate(state, questions)
+
+
+class _SummarizeDecision(beam.DoFn):
+ def process(
+ self,
+ result,
+ timestamp=beam.DoFn.TimestampParam,
+ window=beam.DoFn.WindowParam):
+ answers = result.response.answers
+ yield (
+ result.state,
+ result.response,
+ answers['team'].choice,
+ answers['eligible'].probability >= .5,
+ answers['severity'].score,
+ timestamp.micros, (window.start.micros, window.end.micros))
+
+
+class _LifecycleModel(DecisionModel):
+ def __init__(self):
+ self.enter_count = 0
+ self.exit_calls = []
+ self.active = mock.Mock()
+
+ def __enter__(self):
+ self.enter_count += 1
+ return self.active
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ self.exit_calls.append((exc_type, exc_value, traceback))
+
+ def evaluate(self, state, questions):
+ raise AssertionError('the context returned model must evaluate')
+
+
+def _summarize_failure(
+ failure, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam):
+ return (
+ failure[0],
+ failure[1][1],
+ timestamp.micros, (window.start.micros, window.end.micros))
+
+
+class DecisionTest(unittest.TestCase):
+ def test_dead_letter_output_preserves_failed_state_and_window(self):
+ good = {'id': 'good', 'text': 'bug in invoice', 'priority': 2}
+ bad = {'id': 'bad', 'text': 'refund request', 'priority': 0}
+ with TestPipeline() as pipeline:
+ states = (
+ pipeline
+ | beam.Create([
+ beam.window.TimestampedValue(good, 1),
+ beam.window.TimestampedValue(bad, 11),
+ ])
+ | beam.WindowInto(FixedWindows(10)))
+ decisions, failures = states | EvaluateDecisions(
+ _RejectingModel(), _QUESTIONS).with_exception_handling(
+ exc_class=ValueError)
+ assert_that(
+ decisions | beam.Map(lambda result: result.state),
+ equal_to([good]),
+ label='successful decisions')
+ assert_that(
+ failures | beam.Map(_summarize_failure),
+ equal_to([(
+ bad,
+ "ValueError('Invalid answer probability')",
+ 11_000_000, (10_000_000, 20_000_000))]),
+ label='failed decisions')
+
+ def test_models_share_questions_downstream_and_preserve_window_metadata(self):
+ first = {'id': 'first', 'text': 'bug in invoice', 'priority': 2}
+ second = {'id': 'second', 'text': 'refund request', 'priority': 0}
+ stream = (
+ TestStream().add_elements([
+ beam.window.TimestampedValue(first, 1),
+ beam.window.TimestampedValue(second, 11),
+ ]).advance_watermark_to_infinity())
+
+ expected = {
+ 'keyword': [
+ (
+ first,
+ _response('keyword-v1', 'local-keyword', first, .9, 2),
+ 'support',
+ True,
+ 2.0,
+ 1_000_000, (0, 10_000_000)),
+ (
+ second,
+ _response('keyword-v1', 'local-keyword', second, .2, 0),
+ 'billing',
+ False,
+ 0.0,
+ 11_000_000, (10_000_000, 20_000_000)),
+ ],
+ 'priority': [
+ (
+ first,
+ _response('priority-v1', 'local-priority', first, .8, 2),
+ 'support',
+ True,
+ 2.0,
+ 1_000_000, (0, 10_000_000)),
+ (
+ second,
+ _response('priority-v1', 'local-priority', second, .1, 0),
+ 'billing',
+ False,
+ 0.0,
+ 11_000_000, (10_000_000, 20_000_000)),
+ ],
+ }
+
+ with TestPipeline() as pipeline:
+ states = pipeline | stream | beam.WindowInto(FixedWindows(10))
+ for label, model in (
+ ('keyword', _KeywordModel()),
+ ('priority', _PriorityModel()),
+ ):
+ decisions = states | f'{label} decisions' >> EvaluateDecisions(
+ model, _QUESTIONS)
+ summarized = decisions | f'{label} downstream' >> beam.ParDo(
+ _SummarizeDecision())
+ assert_that(
+ summarized, equal_to(expected[label]), label=f'check {label}')
+
+ def test_context_model_is_reused_across_bundles_and_closed(self):
+ model = _LifecycleModel()
+ model.active.evaluate.side_effect = [
+ _response('test', 'mock', {'id': 'one'}, .9, 1),
+ _response('test', 'mock', {'id': 'two'}, .1, 0),
+ ]
+ dofn = _EvaluateDecisionsDoFn(model, _QUESTIONS)
+
+ dofn.setup()
+ dofn.start_bundle()
+ first = list(dofn.process({'id': 'one'}))
+ dofn.finish_bundle()
+ dofn.start_bundle()
+ second = list(dofn.process({'id': 'two'}))
+ dofn.finish_bundle()
+ dofn.teardown()
+
+ self.assertEqual(1, model.enter_count)
+ self.assertEqual(2, model.active.evaluate.call_count)
+ self.assertEqual(1, len(model.exit_calls))
+ self.assertEqual((None, None, None), model.exit_calls[0])
+ self.assertEqual('one', first[0].state['id'])
+ self.assertEqual('two', second[0].state['id'])
+
+ def test_model_exception_propagates_from_process(self):
+ model = _LifecycleModel()
+ model.active.evaluate.side_effect = RuntimeError('provider failed')
+ dofn = _EvaluateDecisionsDoFn(model, _QUESTIONS)
+
+ dofn.setup()
+ with self.assertRaisesRegex(RuntimeError, 'provider failed'):
+ list(dofn.process({'id': 'failed'}))
+ dofn.teardown()
+ self.assertEqual(1, len(model.exit_calls))
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference.py b/sdks/python/apache_beam/ml/inference/typesafe_inference.py
new file mode 100644
index 000000000000..130a0cafb5e9
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_inference.py
@@ -0,0 +1,164 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Optional TypeSafe adapter for Beam decision models."""
+
+from collections.abc import Mapping
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+
+__all__ = ['JevDecisionModel']
+
+
+def _number(name, field, value, maximum=1):
+ if (isinstance(value, bool) or not isinstance(value, (int, float)) or
+ not 0 <= value <= maximum):
+ raise ValueError(
+ f'Invalid {field} for question {name!r}: '
+ f'expected a finite number in [0, {maximum}]')
+ return value
+
+
+def _probabilities(name, probabilities, keys, key_type):
+ if not isinstance(probabilities, Mapping):
+ raise ValueError(
+ f'Invalid probabilities for question {name!r}: '
+ 'expected a mapping')
+ for key, probability in probabilities.items():
+ if (not isinstance(key, key_type) or isinstance(key, bool) or
+ key not in keys):
+ raise ValueError(
+ f'Invalid probability key for question {name!r}: {key!r}')
+ _number(name, 'probability', probability)
+ return dict(probabilities)
+
+
+def _confidence(name, confidence):
+ return None if confidence is None else _number(name, 'confidence', confidence)
+
+
+def _answer_type(name, answer, expected_type):
+ if not isinstance(answer, expected_type):
+ raise ValueError(
+ f'Invalid answer type for question {name!r}: '
+ f'expected {expected_type.__name__}')
+
+
+class JevDecisionModel(DecisionModel):
+ """Evaluate Beam decision questions with Jev through the TypeSafe SDK.
+
+ Install the optional SDK from ``typesafe_tests_requirements.txt`` and set
+ ``TYPESAFE_API_KEY`` on the workers. The client is created when Beam
+ enters the adapter context; its HTTP pool is reused until teardown.
+ Choice and Score map directly to the SDK primitives. Boolean maps to Noul,
+ whose ``noul`` field becomes :attr:`BooleanAnswer.probability`.
+
+ The SDK retries 429 and server errors with exponential backoff, jitter, and
+ ``Retry-After`` support. Calls and retries block the processing thread; they
+ do not coordinate rate limits across workers. Invalid answer keys, types,
+ labels, numeric ranges, or probability keys raise ``ValueError``.
+
+ Args:
+ model: TypeSafe model identifier, defaulting to ``jev-latest``.
+ timeout: Timeout in seconds for HTTP operations. The SDK handles retries.
+ """
+ def __init__(self, model='jev-latest', timeout=10.0):
+ self.model = model
+ self.timeout = timeout
+ self.client = None
+
+ def __enter__(self):
+ try:
+ from typesafe_sdk import TypeSafeClient
+ except ImportError as error:
+ raise ImportError(
+ 'Install typesafe-sdk to use JevDecisionModel') from error
+ self.client = TypeSafeClient(timeout=self.timeout)
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ self.client.close()
+ self.client = None
+
+ def evaluate(self, state, questions):
+ if self.client is None:
+ raise RuntimeError('Enter JevDecisionModel before evaluating questions')
+ import typesafe_sdk
+
+ sdk_questions = {}
+ for name, question in questions.items():
+ if isinstance(question, ChoiceQuestion):
+ sdk_questions[name] = typesafe_sdk.Choice(
+ instructions=question.instructions, criteria=question.criteria)
+ elif isinstance(question, BooleanQuestion):
+ sdk_questions[name] = typesafe_sdk.Noul(
+ instructions=question.instructions)
+ elif isinstance(question, ScoreQuestion):
+ sdk_questions[name] = typesafe_sdk.Score(
+ instructions=question.instructions, criteria=question.criteria)
+ else:
+ raise TypeError('Unsupported decision question: %r' % type(question))
+
+ try:
+ response = self.client.system_one(
+ state=state, questions=sdk_questions, model=self.model)
+ except typesafe_sdk.TypeSafeAPIResponseValidationError as error:
+ raise ValueError('Invalid Jev response') from error
+ if (not isinstance(response.answers, Mapping) or
+ response.answers.keys() != questions.keys()):
+ raise ValueError(
+ 'Jev answer question names must match the requested '
+ 'question names')
+ answers = {}
+ for name, question in questions.items():
+ answer = response.answers[name]
+ if isinstance(question, ChoiceQuestion):
+ _answer_type(name, answer, typesafe_sdk.ChoiceAnswer)
+ if (not isinstance(answer.choice, str) or
+ answer.choice not in question.criteria):
+ raise ValueError(
+ f'Invalid choice for question {name!r}: '
+ f'{answer.choice!r}')
+ answers[name] = ChoiceAnswer(
+ choice=answer.choice,
+ probabilities=_probabilities(
+ name, answer.probabilities, question.criteria, str),
+ confidence=_confidence(name, answer.confidence))
+ elif isinstance(question, BooleanQuestion):
+ _answer_type(name, answer, typesafe_sdk.NoulAnswer)
+ answers[name] = BooleanAnswer(
+ probability=_number(name, 'probability', answer.noul))
+ else:
+ _answer_type(name, answer, typesafe_sdk.ScoreAnswer)
+ answers[name] = ScoreAnswer(
+ score=_number(
+ name, 'score', answer.score, len(question.criteria) - 1),
+ probabilities=_probabilities(
+ name, answer.probabilities, range(len(question.criteria)), int),
+ confidence=_confidence(name, answer.confidence))
+ return DecisionResponse(
+ answers=answers,
+ model=response.model,
+ provider='typesafe',
+ request_id=response.request_id)
diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py b/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py
new file mode 100644
index 000000000000..acd2234f2ddf
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py
@@ -0,0 +1,93 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Live Jev integration test, enabled with TYPESAFE_API_KEY."""
+
+import os
+import unittest
+
+import pytest
+
+try:
+ import typesafe_sdk
+except ImportError:
+ typesafe_sdk = None
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.test_stream import TestStream
+from apache_beam.testing.util import assert_that
+
+
+@pytest.mark.it_postcommit
+@unittest.skipIf(typesafe_sdk is None, 'typesafe-sdk is not installed')
+@unittest.skipUnless(
+ os.environ.get('TYPESAFE_API_KEY'),
+ 'TYPESAFE_API_KEY environment variable is not set')
+class JevInferenceIT(unittest.TestCase):
+ def test_typed_decisions_in_streaming_pipeline(self):
+ state = {
+ 'message': 'My invoice was charged twice. Please refund the duplicate.'
+ }
+ questions = {
+ 'team': ChoiceQuestion(
+ 'Which team should handle this message?',
+ {
+ 'billing': 'Invoices and refunds',
+ 'technical': 'Bugs and outages',
+ }),
+ 'refund': BooleanQuestion('Does this message ask for a refund?'),
+ 'urgency': ScoreQuestion(
+ 'How urgent is this message?', ('Routine', 'Urgent', 'Emergency')),
+ }
+
+ def check_results(results):
+ assert len(results) == 1, results
+ result = results[0]
+ assert result.state == state
+ assert result.response.provider == 'typesafe'
+ assert result.response.model
+ assert set(result.response.answers) == set(questions)
+ choice = result.response.answers['team']
+ boolean = result.response.answers['refund']
+ score = result.response.answers['urgency']
+ assert isinstance(choice, ChoiceAnswer)
+ assert choice.choice in questions['team'].criteria
+ assert isinstance(boolean, BooleanAnswer)
+ assert 0 <= boolean.probability <= 1
+ assert isinstance(score, ScoreAnswer)
+ assert 0 <= score.score <= 2
+ assert set(score.probabilities) <= {0, 1, 2}
+
+ stream = TestStream().add_elements([state]).advance_watermark_to_infinity()
+ with TestPipeline() as pipeline:
+ results = (
+ pipeline
+ | stream
+ | EvaluateDecisions(JevDecisionModel(), questions))
+ assert_that(results, check_results)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py
new file mode 100644
index 000000000000..5776fc9133e8
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py
@@ -0,0 +1,256 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for TypeSafe's translation to the shared decision types."""
+
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+from parameterized import parameterized
+
+try:
+ import typesafe_sdk
+except ImportError:
+ raise unittest.SkipTest('typesafe-sdk is not installed')
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+
+
+def _choice_answer(**fields):
+ return typesafe_sdk.ChoiceAnswer.model_construct(
+ **{
+ 'choice': 'billing',
+ 'probabilities': {
+ 'billing': 1.0
+ },
+ 'confidence': 0.9,
+ **fields,
+ })
+
+
+def _score_answer(**fields):
+ return typesafe_sdk.ScoreAnswer.model_construct(
+ **{
+ 'score': 0.7,
+ 'probabilities': {
+ 0: 0.3, 1: 0.7
+ },
+ 'confidence': 0.6,
+ **fields,
+ })
+
+
+class JevDecisionModelTest(unittest.TestCase):
+ def test_normalizes_all_primitives_and_reuses_client(self):
+ client = mock.Mock()
+ client.system_one.return_value = SimpleNamespace(
+ answers={
+ 'team': _choice_answer(),
+ 'review': typesafe_sdk.NoulAnswer.model_construct(noul=0.8),
+ 'urgency': _score_answer(),
+ },
+ model='jev-test',
+ request_id='request-1')
+ questions = {
+ 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}),
+ 'review': BooleanQuestion('Does this need review?'),
+ 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']),
+ }
+ state = {'message': 'Review my invoice'}
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient',
+ return_value=client) as factory:
+ adapter = JevDecisionModel(model='jev-test', timeout=2.0)
+ factory.assert_not_called()
+ with adapter:
+ response = adapter.evaluate(state, questions)
+ adapter.evaluate(state, questions)
+ factory.assert_called_once_with(timeout=2.0)
+ client.close.assert_called_once_with()
+ self.assertEqual(2, client.system_one.call_count)
+ request = client.system_one.call_args.kwargs
+ self.assertEqual(state, request['state'])
+ self.assertEqual('jev-test', request['model'])
+ sent = request['questions']
+ self.assertIsInstance(sent['team'], typesafe_sdk.Choice)
+ self.assertEqual(questions['team'].criteria, sent['team'].criteria)
+ self.assertIsInstance(sent['review'], typesafe_sdk.Noul)
+ self.assertEqual(
+ questions['review'].instructions, sent['review'].instructions)
+ self.assertIsInstance(sent['urgency'], typesafe_sdk.Score)
+ self.assertEqual(questions['urgency'].criteria, sent['urgency'].criteria)
+ self.assertEqual(
+ ChoiceAnswer('billing', {'billing': 1.0}, 0.9),
+ response.answers['team'])
+ self.assertEqual(BooleanAnswer(0.8), response.answers['review'])
+ self.assertEqual(
+ ScoreAnswer(0.7, {
+ 0: 0.3, 1: 0.7
+ }, 0.6), response.answers['urgency'])
+ self.assertEqual('jev-test', response.model)
+ self.assertEqual('typesafe', response.provider)
+ self.assertEqual('request-1', response.request_id)
+
+ @parameterized.expand([
+ ('choice_label', 'team', _choice_answer(choice='sales')),
+ ('choice_type', 'team', typesafe_sdk.NoulAnswer.model_construct(noul=.5)),
+ (
+ 'choice_probability_key',
+ 'team',
+ _choice_answer(probabilities={'sales': 1.0})),
+ ('choice_probabilities_type', 'team', _choice_answer(probabilities=[])),
+ (
+ 'choice_probability_range',
+ 'team',
+ _choice_answer(probabilities={'billing': 1.2})),
+ (
+ 'choice_probability_nan',
+ 'team',
+ _choice_answer(probabilities={'billing': float('nan')})),
+ ('choice_confidence', 'team', _choice_answer(confidence=1.3)),
+ ('choice_confidence_type', 'team', _choice_answer(confidence=True)),
+ (
+ 'boolean_range',
+ 'review',
+ typesafe_sdk.NoulAnswer.model_construct(noul=1.4)),
+ (
+ 'boolean_nan',
+ 'review',
+ typesafe_sdk.NoulAnswer.model_construct(noul=float('nan'))),
+ (
+ 'boolean_type',
+ 'review',
+ typesafe_sdk.NoulAnswer.model_construct(noul='0.5')),
+ ('boolean_answer_type', 'review', _choice_answer()),
+ ('score_range', 'urgency', _score_answer(score=3.5)),
+ ('score_negative', 'urgency', _score_answer(score=-.1)),
+ ('score_infinite', 'urgency', _score_answer(score=float('inf'))),
+ ('score_boolean', 'urgency', _score_answer(score=True)),
+ ('score_answer_type', 'urgency', _choice_answer()),
+ (
+ 'score_probability_key',
+ 'urgency',
+ _score_answer(probabilities={999: 1.0})),
+ (
+ 'score_probability_key_type',
+ 'urgency',
+ _score_answer(probabilities={'1': 1.0})),
+ (
+ 'score_probability_fractional_key',
+ 'urgency',
+ _score_answer(probabilities={.5: 1.0})),
+ (
+ 'score_probability_negative',
+ 'urgency',
+ _score_answer(probabilities={0: -.1})),
+ ('score_confidence', 'urgency', _score_answer(confidence=-.1)),
+ ])
+ def test_rejects_malformed_answer(self, unused_name, name, answer):
+ questions = {
+ 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}),
+ 'review': BooleanQuestion('Does this need review?'),
+ 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']),
+ }
+ client = mock.Mock()
+ client.system_one.return_value = SimpleNamespace(
+ answers={name: answer}, model='jev-test', request_id='request-1')
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client):
+ with JevDecisionModel() as adapter:
+ with self.assertRaisesRegex(ValueError, name):
+ adapter.evaluate('Hello', {name: questions[name]})
+
+ @parameterized.expand([
+ ('missing', {}),
+ (
+ 'extra',
+ {
+ 'review': typesafe_sdk.NoulAnswer.model_construct(noul=.5),
+ 'unexpected': typesafe_sdk.NoulAnswer.model_construct(noul=.5),
+ }),
+ ('wrong_container', []),
+ ])
+ def test_rejects_answer_keys(self, unused_name, answers):
+ client = mock.Mock()
+ client.system_one.return_value = SimpleNamespace(
+ answers=answers, model='jev-test', request_id='request-1')
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client):
+ with JevDecisionModel() as adapter:
+ with self.assertRaisesRegex(ValueError, 'question names'):
+ adapter.evaluate(
+ 'Hello', {'review': BooleanQuestion('Needs review?')})
+
+ def test_accepts_boundaries_and_optional_probabilities(self):
+ questions = {
+ 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}),
+ 'review': BooleanQuestion('Does this need review?'),
+ 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']),
+ }
+ client = mock.Mock()
+ for probability in (0, 1):
+ with self.subTest(probability=probability):
+ client.system_one.return_value = SimpleNamespace(
+ answers={
+ 'team': _choice_answer(probabilities={}, confidence=None),
+ 'review': typesafe_sdk.NoulAnswer.model_construct(
+ noul=probability),
+ 'urgency': _score_answer(
+ score=probability, probabilities={}, confidence=None),
+ },
+ model='jev-test',
+ request_id='request-1')
+ with mock.patch.object(typesafe_sdk,
+ 'TypeSafeClient',
+ return_value=client):
+ with JevDecisionModel() as adapter:
+ response = adapter.evaluate('Hello', questions)
+ self.assertEqual(probability, response.answers['review'].probability)
+ self.assertEqual(probability, response.answers['urgency'].score)
+ self.assertIsNone(response.answers['team'].confidence)
+
+ def test_provider_error_propagates_and_client_closes(self):
+ client = mock.Mock()
+ client.system_one.side_effect = TimeoutError('request timed out')
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client):
+ with self.assertRaisesRegex(TimeoutError, 'request timed out'):
+ with JevDecisionModel() as adapter:
+ adapter.evaluate(
+ 'Hello', {'review': BooleanQuestion('Needs review?')})
+ client.close.assert_called_once_with()
+
+ def test_sdk_response_validation_error_is_a_value_error(self):
+ error = typesafe_sdk.TypeSafeAPIResponseValidationError(
+ status=200, body={}, headers={}, field_path='answers.review.noul')
+ client = mock.Mock()
+ client.system_one.side_effect = error
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient', return_value=client):
+ with JevDecisionModel() as adapter:
+ with self.assertRaisesRegex(ValueError,
+ 'Invalid Jev response') as raised:
+ adapter.evaluate(
+ 'Hello', {'review': BooleanQuestion('Needs review?')})
+ self.assertIs(error, raised.exception.__cause__)
+ client.close.assert_called_once_with()
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt b/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt
new file mode 100644
index 000000000000..982081593f5e
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt
@@ -0,0 +1,18 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+typesafe-sdk>=0.7.1,<0.8
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 08c5964898de..a4888617a520 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -613,6 +613,7 @@ def get_portability_package_data():
] + ml_base + qdrant_dependency,
'p312_ml_test': [
'datatable',
+ 'typesafe-sdk>=0.7.1,<0.8',
] + ml_base + qdrant_dependency,
# maintainer: milvus tests only run with this extension. Make sure it
# is covered by docker-in-docker test when changing py version
From 75d7e2c467ff8e074f30cb439502ee6cba1374bf Mon Sep 17 00:00:00 2001
From: Andrew Crites
Date: Mon, 19 Oct 2026 12:42:14 -0700
Subject: [PATCH 044/108] Process deletes before adds when committing
OrderedList state updates. Otherwise we delete the newly added id from the
idTracker. (#40659)
---
.../windmill/state/WindmillOrderedList.java | 28 ++---
.../state/WindmillStateInternalsTest.java | 105 ++++++++++++++++++
2 files changed, 119 insertions(+), 14 deletions(-)
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java
index 03652471a049..0b362d40a666 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillOrderedList.java
@@ -233,6 +233,20 @@ public Windmill.WorkItemCommitRequest persistDirectly(WindmillStateCache.ForKeyA
cleared = false;
}
+ if (!pendingDeletes.isEmpty()) {
+ for (Range range : pendingDeletes.asRanges()) {
+ Windmill.TagSortedListDeleteRequest.Builder deletesBuilder =
+ updatesBuilder.addDeletesBuilder();
+ deletesBuilder.setRange(
+ Windmill.SortedListRange.newBuilder()
+ .setStart(WindmillTimeUtils.harnessToWindmillTimestamp(range.lowerEndpoint()))
+ .setLimit(WindmillTimeUtils.harnessToWindmillTimestamp(range.upperEndpoint())));
+ deletesBuilder.build();
+ idTracker.remove(range);
+ }
+ pendingDeletes.clear();
+ }
+
if (!pendingAdds.isEmpty()) {
// TODO(reuvenlax): Once we start caching data, we should remove this line. We have it
// here now
@@ -261,20 +275,6 @@ public Windmill.WorkItemCommitRequest persistDirectly(WindmillStateCache.ForKeyA
pendingAdds.clear();
insertBuilder.build();
}
-
- if (!pendingDeletes.isEmpty()) {
- for (Range range : pendingDeletes.asRanges()) {
- Windmill.TagSortedListDeleteRequest.Builder deletesBuilder =
- updatesBuilder.addDeletesBuilder();
- deletesBuilder.setRange(
- Windmill.SortedListRange.newBuilder()
- .setStart(WindmillTimeUtils.harnessToWindmillTimestamp(range.lowerEndpoint()))
- .setLimit(WindmillTimeUtils.harnessToWindmillTimestamp(range.upperEndpoint())));
- deletesBuilder.build();
- idTracker.remove(range);
- }
- pendingDeletes.clear();
- }
} catch (ExecutionException | InterruptedException e) {
throw new RuntimeException(e);
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java
index 87b746089f11..7a68974f4128 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java
@@ -2386,6 +2386,111 @@ public void testOrderedListDeleteRangePersist() {
assertBuildable(commitBuilder);
}
+ @Test
+ public void testOrderedListDeleteRangeThenAddPersist() throws Exception {
+ SettableFuture, RangeSet>> orderedListFuture = SettableFuture.create();
+ orderedListFuture.set(null);
+ SettableFuture, RangeSet>> deletionsFuture =
+ SettableFuture.create();
+ deletionsFuture.set(null);
+ when(mockReader.valueFuture(
+ systemKey(NAMESPACE, "orderedList" + IdTracker.IDS_AVAILABLE_STR),
+ STATE_FAMILY,
+ IdTracker.IDS_AVAILABLE_CODER))
+ .thenReturn(orderedListFuture);
+ when(mockReader.valueFuture(
+ systemKey(NAMESPACE, "orderedList" + IdTracker.DELETIONS_STR),
+ STATE_FAMILY,
+ IdTracker.SUBRANGE_DELETIONS_CODER))
+ .thenReturn(deletionsFuture);
+
+ StateTag> addr =
+ StateTags.orderedList("orderedList", StringUtf8Coder.of());
+ OrderedListState orderedListState = underTest.state(NAMESPACE, addr);
+
+ // Clear a sub-range [2ms, 4ms) and then add an element at 2ms within that cleared range.
+ orderedListState.clearRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4));
+ orderedListState.add(TimestampedValue.of("hello", Instant.ofEpochMilli(2)));
+
+ Windmill.WorkItemCommitRequest.Builder commitBuilder =
+ Windmill.WorkItemCommitRequest.newBuilder();
+ underTest.persist(commitBuilder);
+
+ assertEquals(1, commitBuilder.getSortedListUpdatesCount());
+ TagSortedListUpdateRequest updates = commitBuilder.getSortedListUpdates(0);
+ assertEquals(1, updates.getDeletesCount());
+ assertEquals(2000, updates.getDeletes(0).getRange().getStart());
+ assertEquals(4000, updates.getDeletes(0).getRange().getLimit());
+ assertEquals(1, updates.getInsertsCount());
+ assertEquals(1, updates.getInserts(0).getEntriesCount());
+ assertEquals("hello", updates.getInserts(0).getEntries(0).getValue().toStringUtf8());
+ assertEquals(2000, updates.getInserts(0).getEntries(0).getSortKey());
+ assertEquals(IdTracker.NEW_RANGE_MIN_ID, updates.getInserts(0).getEntries(0).getId());
+ assertBuildable(commitBuilder);
+
+ // Verify that the sub-range deletion for [2ms, 3ms) was removed when "hello" was added at 2ms,
+ // so clearing the rest of the 5-minute bucket [0, 2ms) and [4ms, 5m) in the next work item
+ // does not wipe out id tracking for the bucket.
+ resetUnderTest();
+ orderedListState = underTest.state(NAMESPACE, addr);
+ orderedListState.clearRange(Instant.EPOCH, Instant.ofEpochMilli(2));
+ orderedListState.clearRange(Instant.ofEpochMilli(4), Instant.ofEpochSecond(300));
+
+ commitBuilder = Windmill.WorkItemCommitRequest.newBuilder();
+ underTest.persist(commitBuilder);
+ assertBuildable(commitBuilder);
+
+ // Adding another element at 2ms in a subsequent work item should get the next id (1), not 0.
+ resetUnderTest();
+ orderedListState = underTest.state(NAMESPACE, addr);
+ orderedListState.add(TimestampedValue.of("world", Instant.ofEpochMilli(2)));
+
+ commitBuilder = Windmill.WorkItemCommitRequest.newBuilder();
+ underTest.persist(commitBuilder);
+
+ updates = commitBuilder.getSortedListUpdates(0);
+ assertEquals(1, updates.getInsertsCount());
+ assertEquals(1, updates.getInserts(0).getEntriesCount());
+ assertEquals("world", updates.getInserts(0).getEntries(0).getValue().toStringUtf8());
+ assertEquals(2000, updates.getInserts(0).getEntries(0).getSortKey());
+ assertEquals(IdTracker.NEW_RANGE_MIN_ID + 1, updates.getInserts(0).getEntries(0).getId());
+ assertBuildable(commitBuilder);
+
+ // Now clear the entire 5-minute bucket [0, 5m) and add a new element in the same work item.
+ // The delete should be processed before the add, resetting the bucket's ids to start at
+ // NEW_RANGE_MIN_ID (0) for the new element while preserving its id tracking for future adds.
+ resetUnderTest();
+ orderedListState = underTest.state(NAMESPACE, addr);
+ orderedListState.clearRange(Instant.EPOCH, Instant.ofEpochSecond(300));
+ orderedListState.add(TimestampedValue.of("afterFullClear1", Instant.ofEpochMilli(2)));
+
+ commitBuilder = Windmill.WorkItemCommitRequest.newBuilder();
+ underTest.persist(commitBuilder);
+
+ updates = commitBuilder.getSortedListUpdates(0);
+ assertEquals(1, updates.getDeletesCount());
+ assertEquals(1, updates.getInsertsCount());
+ assertEquals(1, updates.getInserts(0).getEntriesCount());
+ assertEquals("afterFullClear1", updates.getInserts(0).getEntries(0).getValue().toStringUtf8());
+ assertEquals(IdTracker.NEW_RANGE_MIN_ID, updates.getInserts(0).getEntries(0).getId());
+ assertBuildable(commitBuilder);
+
+ // A subsequent add in the same bucket should receive NEW_RANGE_MIN_ID + 1.
+ resetUnderTest();
+ orderedListState = underTest.state(NAMESPACE, addr);
+ orderedListState.add(TimestampedValue.of("afterFullClear2", Instant.ofEpochMilli(2)));
+
+ commitBuilder = Windmill.WorkItemCommitRequest.newBuilder();
+ underTest.persist(commitBuilder);
+
+ updates = commitBuilder.getSortedListUpdates(0);
+ assertEquals(1, updates.getInsertsCount());
+ assertEquals(1, updates.getInserts(0).getEntriesCount());
+ assertEquals("afterFullClear2", updates.getInserts(0).getEntries(0).getValue().toStringUtf8());
+ assertEquals(IdTracker.NEW_RANGE_MIN_ID + 1, updates.getInserts(0).getEntries(0).getId());
+ assertBuildable(commitBuilder);
+ }
+
@Test
public void testOrderedListMergePendingAdds() {
SettableFuture, RangeSet>> orderedListFuture = SettableFuture.create();
From b4ecb26ada0b8355c8a32ddb0799ae447339c8cf Mon Sep 17 00:00:00 2001
From: CH-Abhinav