[pull] master from apache:master

.github/workflows/build_release_candidate.yml

@@ -295,7 +295,7 @@ jobs:
         # settings.xml file
         run: rm ~/.m2/settings.xml || true
       - name: Login to Docker Hub
-        uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4
+        uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0
         with:
           username: ${{ secrets.DOCKERHUB_USER }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}

.github/workflows/finalize_release.yml

@@ -44,7 +44,7 @@ jobs:
     runs-on: [self-hosted, ubuntu-24.04, main]
     steps:
     - name: Login to Docker Hub
-      uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4
+      uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0
       with:
         username: ${{ secrets.DOCKERHUB_USER }}
         password: ${{ secrets.DOCKERHUB_TOKEN }}

sdks/python/apache_beam/runners/worker/bundle_processor.py

@@ -641,6 +641,8 @@ def __init__(
     self._value_coder = value_coder
     self._cleared = False
     self._added_elements: set[Any] = set()
+    # Track outstanding async state requests to await them at commit time.
+    self._futures = []
 
   def _compact_data(self, rewrite=True):
     accumulator = set(
@@ -650,9 +652,11 @@ def _compact_data(self, rewrite=True):
             self._added_elements))
 
     if rewrite and accumulator:
-      self._state_handler.clear(self._state_key)
-      self._state_handler.extend(
-          self._state_key, self._value_coder.get_impl(), accumulator)
+      # Compaction writes are asynchronous; queue them so they are not lost.
+      self._futures.append(self._state_handler.clear(self._state_key))
+      self._futures.append(
+          self._state_handler.extend(
+              self._state_key, self._value_coder.get_impl(), accumulator))
 
       # Since everthing is already committed so we can safely reinitialize
       # added_elements here.
@@ -666,7 +670,7 @@ def read(self) -> set[Any]:
   def add(self, value: Any) -> None:
     if self._cleared:
       # This is a good time explicitly clear.
-      self._state_handler.clear(self._state_key)
+      self._futures.append(self._state_handler.clear(self._state_key))
       self._cleared = False
 
     self._added_elements.add(value)
@@ -678,15 +682,26 @@ def clear(self) -> None:
     self._added_elements = set()
 
   def commit(self) -> None:
-    to_await = None
     if self._cleared:
-      to_await = self._state_handler.clear(self._state_key)
+      self._futures.append(self._state_handler.clear(self._state_key))
+      self._cleared = False
     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()
+      self._futures.append(
+          self._state_handler.extend(
+              self._state_key,
+              self._value_coder.get_impl(),
+              self._added_elements))
+      self._added_elements = set()
+
+    # Block on all outstanding async state requests to ensure data is committed.
+    # We must swap and clear self._futures before awaiting them. Awaiting a future
+    # yields control, during which new futures could be appended to self._futures.
+    all_futures = self._futures
+    self._futures = []
+
+    for f in all_futures:
+      if f:
+        f.get()
 
 
 class RangeSet:

sdks/python/apache_beam/testing/benchmarks/cloudml/criteo_tft/criteo.py

@@ -38,6 +38,21 @@ def get_transformed_categorical_column_name(column_name_or_id):
   return column_name + '_id'
 
 
+def fill_in_missing(feature, default_value):
+  """Fills missing values in a rank 2 SparseTensor.
+
+  Args:
+    feature: A rank 2 SparseTensor with at most one value per row.
+    default_value: The value to fill in for missing entries.
+
+  Returns:
+    A rank 1 Tensor with missing entries filled in.
+  """
+  feature = tft.sparse_tensor_to_dense_with_shape(
+      feature, [None, 1], default_value=default_value)
+  return tf.squeeze(feature, axis=1)
+
+
 _INTEGER_COLUMN_NAMES = [
     'int-feature-{}'.format(column_idx) for column_idx in range(1, 14)
 ]
@@ -132,23 +147,12 @@ def preprocessing_fn(inputs):
     result = {'clicked': inputs['clicked']}
     for name in _INTEGER_COLUMN_NAMES:
       feature = inputs[name]
-      # TODO(https://github.com/apache/beam/issues/24902):
-      #  Replace this boilerplate with a helper function.
-      # This is a SparseTensor because it is optional. Here we fill in a
-      # default value when it is missing.
-      feature = tft.sparse_tensor_to_dense_with_shape(
-          feature, [None, 1], default_value=-1)
-      # Reshaping from a batch of vectors of size 1 to a batch of scalars and
-      # adding a bucketized version.
-      feature = tf.squeeze(feature, axis=1)
+      feature = fill_in_missing(feature, -1)
       result[name] = feature
       result[name + '_bucketized'] = tft.bucketize(feature, _NUM_BUCKETS)
     for name in _CATEGORICAL_COLUMN_NAMES:
       feature = inputs[name]
-      # Similar to for integer columns, but use '' as default.
-      feature = tft.sparse_tensor_to_dense_with_shape(
-          feature, [None, 1], default_value='')
-      feature = tf.squeeze(feature, axis=1)
+      feature = fill_in_missing(feature, '')
       result[get_transformed_categorical_column_name(
           name)] = tft.compute_and_apply_vocabulary(
               feature, frequency_threshold=frequency_threshold)

sdks/python/apache_beam/testing/benchmarks/cloudml/criteo_tft/criteo_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.
+#
+
+import unittest
+
+try:
+  import tensorflow as tf
+
+  from apache_beam.testing.benchmarks.cloudml.criteo_tft import criteo
+except ImportError:
+  raise unittest.SkipTest('Dependencies are not installed')
+
+
+class CriteoTest(tf.test.TestCase):
+  def test_fill_in_missing_int_feature(self):
+    feature = tf.SparseTensor(
+        indices=[[0, 0], [2, 0]],
+        values=tf.constant([10, 30], dtype=tf.int64),
+        dense_shape=[3, 1])
+
+    result = criteo.fill_in_missing(feature, -1)
+
+    self.assertAllEqual(result, [10, -1, 30])
+    self.assertEqual(result.shape.rank, 1)
+
+  def test_fill_in_missing_int_feature_traces_with_dynamic_shape(self):
+    @tf.function(
+        input_signature=[
+            tf.SparseTensorSpec(shape=[None, None], dtype=tf.int64)
+        ])
+    def fill_in_missing(feature):
+      return criteo.fill_in_missing(feature, -1)
+
+    feature = tf.SparseTensor(
+        indices=[[0, 0], [2, 0]],
+        values=tf.constant([10, 30], dtype=tf.int64),
+        dense_shape=[3, 1])
+
+    result = fill_in_missing(feature)
+
+    self.assertAllEqual(result, [10, -1, 30])
+    self.assertEqual(result.shape.rank, 1)
+
+  def test_fill_in_missing_all_missing_int_feature(self):
+    feature = tf.SparseTensor(
+        indices=tf.zeros([0, 2], dtype=tf.int64),
+        values=tf.constant([], dtype=tf.int64),
+        dense_shape=[3, 0])
+
+    result = criteo.fill_in_missing(feature, -1)
+
+    self.assertAllEqual(result, [-1, -1, -1])
+    self.assertEqual(result.shape.rank, 1)
+
+  def test_fill_in_missing_string_feature(self):
+    feature = tf.SparseTensor(
+        indices=[[0, 0], [2, 0]],
+        values=tf.constant(['a', 'c'], dtype=tf.string),
+        dense_shape=[3, 1])
+
+    result = criteo.fill_in_missing(feature, '')
+
+    self.assertAllEqual(result, [b'a', b'', b'c'])
+    self.assertEqual(result.shape.rank, 1)
+
+  def test_fill_in_missing_all_missing_string_feature(self):
+    feature = tf.SparseTensor(
+        indices=tf.zeros([0, 2], dtype=tf.int64),
+        values=tf.constant([], dtype=tf.string),
+        dense_shape=[3, 0])
+
+    result = criteo.fill_in_missing(feature, '')
+
+    self.assertAllEqual(result, [b'', b'', b''])
+    self.assertEqual(result.shape.rank, 1)
+
+
+if __name__ == '__main__':
+  unittest.main()

sdks/python/apache_beam/transforms/userstate_test.py

@@ -18,6 +18,9 @@
 """Unit tests for the Beam State and Timer API interfaces."""
 # pytype: skip-file
 
+import queue
+import threading
+import time
 import unittest
 from typing import Any
 
@@ -34,6 +37,8 @@
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.runners import pipeline_context
 from apache_beam.runners.common import DoFnSignature
+from apache_beam.runners.worker.sdk_worker import GrpcStateHandler
+from apache_beam.runners.worker.sdk_worker import _Future
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_stream import TestStream
 from apache_beam.testing.util import assert_that
@@ -719,6 +724,70 @@ def process(self, element, set_state=beam.DoFn.StateParam(SET_STATE)):
       actual_values = (values | beam.ParDo(SetStatefulDoFn()))
       assert_that(actual_values, equal_to([1, 3, 6, 10, 10]))
 
+  # Mock random to always return 1.0 to force compaction on every add.
+  @mock.patch(
+      'apache_beam.runners.worker.bundle_processor.random.random', lambda: 1.0)
+  def test_stateful_set_state_compaction_race_portably(self):
+    old_request = GrpcStateHandler._request
+    request_queue = queue.Queue()
+
+    def worker():
+      while True:
+        handler, request, future, instruction_id = request_queue.get()
+        time.sleep(0.1)  # Simulate latency for each request sequentially.
+        handler._context.process_instruction_id = instruction_id
+        underlying_future = old_request(handler, request)
+        underlying_future.wait()
+        future.set(underlying_future.get())
+        request_queue.task_done()
+
+    t = threading.Thread(target=worker, daemon=True)
+    t.start()
+
+    def delayed_request(self, request):
+      if request.HasField('append') or request.HasField('clear'):
+        future = _Future()
+        instruction_id = getattr(self._context, 'process_instruction_id', None)
+        request_queue.put((self, request, future, instruction_id))
+        return future
+      else:
+        return old_request(self, request)
+
+    GrpcStateHandler._request = delayed_request
+
+    class SetStatefulDoFn(beam.DoFn):
+
+      SET_STATE = SetStateSpec('buffer', VarIntCoder())
+
+      def process(self, element, set_state=beam.DoFn.StateParam(SET_STATE)):
+        _, value = element
+        aggregated_value = 0
+        set_state.add(value)
+        for saved_value in set_state.read():
+          aggregated_value += saved_value
+        yield aggregated_value
+
+    try:
+      options = PipelineOptions([
+          '--max_cache_memory_usage_mb=100',
+          '--environment_type=LOOPBACK',
+      ])
+      with TestPipeline(options=options) as p:
+        test_stream = (
+            TestStream(
+                coder=beam.coders.TupleCoder((
+                    beam.coders.StrUtf8Coder(), beam.coders.VarIntCoder()
+                ))).advance_watermark_to(10).add_elements([
+                    ('key', 1)
+                ]).advance_watermark_to(20).add_elements([
+                    ('key', 2)
+                ]).advance_watermark_to(30))
+        actual_values = (p | test_stream | beam.ParDo(SetStatefulDoFn()))
+        assert_that(actual_values, equal_to([1, 3]))
+
+    finally:
+      GrpcStateHandler._request = old_request
+
   def test_stateful_set_state_clean_portably(self):
     class SetStateClearingStatefulDoFn(beam.DoFn):