apache-beam 2.70.0rc4__tar.gz → 2.71.0__tar.gz

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-beam
-Version: 2.70.0rc4
+Version: 2.71.0
 Summary: Apache Beam SDK for Python
 Home-page: https://beam.apache.org
 Download-URL: https://pypi.python.org/pypi/apache-beam
@@ -20,13 +20,15 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: cryptography<48.0.0,>=39.0.0
+Requires-Dist: envoy-data-plane<2,>=1.0.3; python_version >= "3.13"
+Requires-Dist: envoy-data-plane<0.3.0; python_version < "3.13"
 Requires-Dist: fastavro<2,>=0.23.6
 Requires-Dist: fasteners<1.0,>=0.3
 Requires-Dist: grpcio!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,<1.66.0,<2,>=1.33.1; python_version <= "3.12"
 Requires-Dist: grpcio>=1.67.0; python_version >= "3.13"
 Requires-Dist: httplib2<0.23.0,>=0.8
 Requires-Dist: jsonpickle<4.0.0,>=3.0.0
-Requires-Dist: numpy<2.3.0,>=1.14.3
+Requires-Dist: numpy<2.5.0,>=1.14.3
 Requires-Dist: objsize<0.8.0,>=0.6.1
 Requires-Dist: packaging>=22.0
 Requires-Dist: pymongo<5.0.0,>=3.8.0
@@ -39,7 +41,7 @@ Requires-Dist: sortedcontainers>=2.4.0
 Requires-Dist: typing-extensions>=3.7.0
 Requires-Dist: zstandard<1,>=0.18.0
 Requires-Dist: pyyaml<7.0.0,>=3.12
-Requires-Dist: beartype<0.22.0,>=0.21.0
+Requires-Dist: beartype<0.23.0,>=0.21.0
 Requires-Dist: pyarrow<19.0.0,>=3.0.0
 Requires-Dist: pyarrow-hotfix<1
 Provides-Extra: dill

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/apache_beam/coders/coder_impl.py

@@ -32,6 +32,7 @@ For internal use only; no backwards-compatibility guarantees.
 
 import decimal
 import enum
+import functools
 import itertools
 import json
 import logging
@@ -354,6 +355,7 @@ DATACLASS_TYPE = 101
 NAMED_TUPLE_TYPE = 102
 ENUM_TYPE = 103
 NESTED_STATE_TYPE = 104
+DATACLASS_KW_ONLY_TYPE = 105
 
 # Types that can be encoded as iterables, but are not literally
 # lists, etc. due to being lazy.  The actual type is not preserved
@@ -374,6 +376,18 @@ def _verify_dill_compat():
     raise RuntimeError(base_error + f". Found dill version '{dill.__version__}")
 
 
+dataclass_uses_kw_only: Callable[[Any], bool]
+if dataclasses:
+  # Cache the result to avoid multiple checks for the same dataclass type.
+  @functools.cache
+  def dataclass_uses_kw_only(cls) -> bool:
+    return any(
+        field.init and field.kw_only for field in dataclasses.fields(cls))
+
+else:
+  dataclass_uses_kw_only = lambda cls: False
+
+
 class FastPrimitivesCoderImpl(StreamCoderImpl):
   """For internal use only; no backwards-compatibility guarantees."""
   def __init__(
@@ -497,18 +511,25 @@ class FastPrimitivesCoderImpl(StreamCoderImpl):
       self.encode_type(type(value), stream)
       stream.write(value.SerializePartialToString(deterministic=True), True)
     elif dataclasses and dataclasses.is_dataclass(value):
-      stream.write_byte(DATACLASS_TYPE)
       if not type(value).__dataclass_params__.frozen:
         raise TypeError(
             "Unable to deterministically encode non-frozen '%s' of type '%s' "
             "for the input of '%s'" %
             (value, type(value), self.requires_deterministic_step_label))
-      self.encode_type(type(value), stream)
-      values = [
-          getattr(value, field.name) for field in dataclasses.fields(value)
-      ]
+      init_fields = [field for field in dataclasses.fields(value) if field.init]
       try:
-        self.iterable_coder_impl.encode_to_stream(values, stream, True)
+        if dataclass_uses_kw_only(type(value)):
+          stream.write_byte(DATACLASS_KW_ONLY_TYPE)
+          self.encode_type(type(value), stream)
+          stream.write_var_int64(len(init_fields))
+          for field in init_fields:
+            stream.write(field.name.encode("utf-8"), True)
+            self.encode_to_stream(getattr(value, field.name), stream, True)
+        else:  # Not using kw_only, we can pass parameters by position.
+          stream.write_byte(DATACLASS_TYPE)
+          self.encode_type(type(value), stream)
+          values = [getattr(value, field.name) for field in init_fields]
+          self.iterable_coder_impl.encode_to_stream(values, stream, True)
       except Exception as e:
         raise TypeError(self._deterministic_encoding_error_msg(value)) from e
     elif isinstance(value, tuple) and hasattr(type(value), '_fields'):
@@ -616,6 +637,14 @@ class FastPrimitivesCoderImpl(StreamCoderImpl):
       msg = cls()
       msg.ParseFromString(stream.read_all(True))
       return msg
+    elif t == DATACLASS_KW_ONLY_TYPE:
+      cls = self.decode_type(stream)
+      vlen = stream.read_var_int64()
+      fields = {}
+      for _ in range(vlen):
+        field_name = stream.read_all(True).decode('utf-8')
+        fields[field_name] = self.decode_from_stream(stream, True)
+      return cls(**fields)
     elif t == DATACLASS_TYPE or t == NAMED_TUPLE_TYPE:
       cls = self.decode_type(stream)
       return cls(*self.iterable_coder_impl.decode_from_stream(stream, True))

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/apache_beam/coders/coders_test_common.py

@@ -113,6 +113,11 @@ if dataclasses is not None:
     a: Any
     b: int
 
+  @dataclasses.dataclass(frozen=True, kw_only=True)
+  class FrozenKwOnlyDataClass:
+    c: int
+    d: int
+
   @dataclasses.dataclass
   class UnFrozenDataClass:
     x: int
@@ -303,9 +308,11 @@ class CodersTest(unittest.TestCase):
 
     if dataclasses is not None:
       self.check_coder(deterministic_coder, FrozenDataClass(1, 2))
+      self.check_coder(deterministic_coder, FrozenKwOnlyDataClass(c=1, d=2))
 
       with self.assertRaises(TypeError):
         self.check_coder(deterministic_coder, UnFrozenDataClass(1, 2))
+
       with self.assertRaises(TypeError):
         self.check_coder(
             deterministic_coder, FrozenDataClass(UnFrozenDataClass(1, 2), 3))
@@ -742,6 +749,7 @@ class CodersTest(unittest.TestCase):
         from apache_beam.coders.coders_test_common import DefinesGetState
         from apache_beam.coders.coders_test_common import DefinesGetAndSetState
         from apache_beam.coders.coders_test_common import FrozenDataClass
+        from apache_beam.coders.coders_test_common import FrozenKwOnlyDataClass
 
 
         from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message
@@ -777,6 +785,8 @@ class CodersTest(unittest.TestCase):
         test_cases.extend([
             ("frozen_dataclass", FrozenDataClass(1, 2)),
             ("frozen_dataclass_list", [FrozenDataClass(1, 2), FrozenDataClass(3, 4)]),
+            ("frozen_kwonly_dataclass", FrozenKwOnlyDataClass(c=1, d=2)),
+            ("frozen_kwonly_dataclass_list", [FrozenKwOnlyDataClass(c=1, d=2), FrozenKwOnlyDataClass(c=3, d=4)]),
         ])
 
         compat_version = {'"'+ compat_version +'"' if compat_version else None}

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/apache_beam/coders/typecoders.py

@@ -114,6 +114,14 @@ class CoderRegistry(object):
       typehint_coder_class: Type[coders.Coder]) -> None:
     self._coders[typehint_type] = typehint_coder_class
 
+  @staticmethod
+  def _normalize_typehint_type(typehint_type):
+    if typehint_type.__module__ == '__main__':
+      # See https://github.com/apache/beam/issues/21541
+      # TODO(robertwb): Remove once all runners are portable.
+      return getattr(typehint_type, '__name__', str(typehint_type))
+    return typehint_type
+
   def register_coder(
       self, typehint_type: Any,
       typehint_coder_class: Type[coders.Coder]) -> None:
@@ -123,11 +131,8 @@ class CoderRegistry(object):
           'Received %r instead.' % typehint_coder_class)
     if typehint_type not in self.custom_types:
       self.custom_types.append(typehint_type)
-    if typehint_type.__module__ == '__main__':
-      # See https://github.com/apache/beam/issues/21541
-      # TODO(robertwb): Remove once all runners are portable.
-      typehint_type = getattr(typehint_type, '__name__', str(typehint_type))
-    self._register_coder_internal(typehint_type, typehint_coder_class)
+    self._register_coder_internal(
+        self._normalize_typehint_type(typehint_type), typehint_coder_class)
 
   def get_coder(self, typehint: Any) -> coders.Coder:
     if typehint and typehint.__module__ == '__main__':
@@ -170,9 +175,15 @@ class CoderRegistry(object):
       coder = self._fallback_coder
     return coder.from_type_hint(typehint, self)
 
-  def get_custom_type_coder_tuples(self, types):
+  def get_custom_type_coder_tuples(self, types=None):
     """Returns type/coder tuples for all custom types passed in."""
-    return [(t, self._coders[t]) for t in types if t in self.custom_types]
+    return [(t, self._coders[self._normalize_typehint_type(t)])
+            for t in self.custom_types if (types is None or t in types)]
+
+  def load_custom_type_coder_tuples(self, type_coder):
+    """Load type/coder tuples into coder registry."""
+    for t, c in type_coder:
+      self.register_coder(t, c)
 
   def verify_deterministic(self, key_coder, op_name, silent=True):
     if not key_coder.is_deterministic():

apache_beam-2.71.0/apache_beam/examples/inference/rate_limiter_vertex_ai.py

@@ -0,0 +1,85 @@
+#
+# 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 simple example demonstrating usage of the EnvoyRateLimiter with Vertex AI.
+"""
+
+import argparse
+import logging
+
+import apache_beam as beam
+from apache_beam.io.components.rate_limiter import EnvoyRateLimiter
+from apache_beam.ml.inference.base import RunInference
+from apache_beam.ml.inference.vertex_ai_inference import VertexAIModelHandlerJSON
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import SetupOptions
+
+
+def run(argv=None):
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--project',
+      dest='project',
+      help='The Google Cloud project ID for Vertex AI.')
+  parser.add_argument(
+      '--location',
+      dest='location',
+      help='The Google Cloud location (e.g. us-central1) for Vertex AI.')
+  parser.add_argument(
+      '--endpoint_id',
+      dest='endpoint_id',
+      help='The ID of the Vertex AI endpoint.')
+  parser.add_argument(
+      '--rls_address',
+      dest='rls_address',
+      help='The address of the Envoy Rate Limit Service (e.g. localhost:8081).')
+
+  known_args, pipeline_args = parser.parse_known_args(argv)
+  pipeline_options = PipelineOptions(pipeline_args)
+  pipeline_options.view_as(SetupOptions).save_main_session = True
+
+  # Initialize the EnvoyRateLimiter
+  rate_limiter = EnvoyRateLimiter(
+      service_address=known_args.rls_address,
+      domain="mongo_cps",
+      descriptors=[{
+          "database": "users"
+      }],
+      namespace='example_pipeline')
+
+  # Initialize the VertexAIModelHandler with the rate limiter
+  model_handler = VertexAIModelHandlerJSON(
+      endpoint_id=known_args.endpoint_id,
+      project=known_args.project,
+      location=known_args.location,
+      rate_limiter=rate_limiter)
+
+  # Input features for the model
+  features = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0],
+              [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]]
+
+  with beam.Pipeline(options=pipeline_options) as p:
+    _ = (
+        p
+        | 'CreateInputs' >> beam.Create(features)
+        | 'RunInference' >> RunInference(model_handler)
+        | 'PrintPredictions' >> beam.Map(logging.info))
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  run()

apache_beam-2.71.0/apache_beam/examples/rate_limiter_simple.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.
+#
+
+"""A simple example demonstrating usage of the EnvoyRateLimiter in a Beam
+pipeline.
+"""
+
+import argparse
+import logging
+import time
+
+import apache_beam as beam
+from apache_beam.io.components.rate_limiter import EnvoyRateLimiter
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.utils import shared
+
+
+class SampleApiDoFn(beam.DoFn):
+  """A DoFn that simulates calling an external API with rate limiting."""
+  def __init__(self, rls_address, domain, descriptors):
+    self.rls_address = rls_address
+    self.domain = domain
+    self.descriptors = descriptors
+    self._shared = shared.Shared()
+    self.rate_limiter = None
+
+  def setup(self):
+    # Initialize the rate limiter in setup()
+    # We use shared.Shared() to ensure only one RateLimiter instance is created
+    # per worker and shared across threads.
+    def init_limiter():
+      logging.info("Connecting to Envoy RLS at %s", self.rls_address)
+      return EnvoyRateLimiter(
+          service_address=self.rls_address,
+          domain=self.domain,
+          descriptors=self.descriptors,
+          namespace='example_pipeline')
+
+    self.rate_limiter = self._shared.acquire(init_limiter)
+
+  def process(self, element):
+    self.rate_limiter.allow()
+
+    # Process the element mock API call
+    logging.info("Processing element: %s", element)
+    time.sleep(0.1)
+    yield element
+
+
+def parse_known_args(argv):
+  """Parses args for the workflow."""
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--rls_address',
+      default='localhost:8081',
+      help='Address of the Envoy Rate Limit Service')
+  return parser.parse_known_args(argv)
+
+
+def run(argv=None):
+  known_args, pipeline_args = parse_known_args(argv)
+  pipeline_options = PipelineOptions(pipeline_args)
+
+  with beam.Pipeline(options=pipeline_options) as p:
+    _ = (
+        p
+        | 'Create' >> beam.Create(range(100))
+        | 'RateLimit' >> beam.ParDo(
+            SampleApiDoFn(
+                rls_address=known_args.rls_address,
+                domain="mongo_cps",
+                descriptors=[{
+                    "database": "users"
+                }])))
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  run()

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py

@@ -177,7 +177,7 @@ class EnrichmentTest(unittest.TestCase):
         enrichment_with_google_cloudsql_pg()
         output = mock_stdout.getvalue().splitlines()
         expected = validate_enrichment_with_google_cloudsql_pg()
-        self.assertEqual(output, expected)
+        self.assertEqual(sorted(output), sorted(expected))
     except Exception as e:
       self.fail(f"Test failed with unexpected error: {e}")
 
@@ -188,7 +188,7 @@ class EnrichmentTest(unittest.TestCase):
         enrichment_with_external_pg()
         output = mock_stdout.getvalue().splitlines()
         expected = validate_enrichment_with_external_pg()
-        self.assertEqual(output, expected)
+        self.assertEqual(sorted(output), sorted(expected))
     except (TestContainerStartupError, TestContainerTeardownError) as e:
       raise unittest.SkipTest(str(e))
     except Exception as e:
@@ -201,7 +201,7 @@ class EnrichmentTest(unittest.TestCase):
         enrichment_with_external_mysql()
         output = mock_stdout.getvalue().splitlines()
         expected = validate_enrichment_with_external_mysql()
-        self.assertEqual(output, expected)
+        self.assertEqual(sorted(output), sorted(expected))
     except (TestContainerStartupError, TestContainerTeardownError) as e:
       raise unittest.SkipTest(str(e))
     except Exception as e:
@@ -214,7 +214,7 @@ class EnrichmentTest(unittest.TestCase):
         enrichment_with_external_sqlserver()
         output = mock_stdout.getvalue().splitlines()
         expected = validate_enrichment_with_external_sqlserver()
-        self.assertEqual(output, expected)
+        self.assertEqual(sorted(output), sorted(expected))
     except (TestContainerStartupError, TestContainerTeardownError) as e:
       raise unittest.SkipTest(str(e))
     except Exception as e:

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/apache_beam/internal/cloudpickle_pickler.py

@@ -252,12 +252,35 @@ def _lock_reducer(obj):
 
 
 def dump_session(file_path):
-  # It is possible to dump session with cloudpickle. However, since references
-  # are saved it should not be necessary. See https://s.apache.org/beam-picklers
-  pass
+  # Since References are saved (https://s.apache.org/beam-picklers), we only
+  # dump supported Beam Registries (currently only logical type registry)
+  from apache_beam.coders import typecoders
+  from apache_beam.typehints import schemas
+
+  with _pickle_lock, open(file_path, 'wb') as file:
+    coder_reg = typecoders.registry.get_custom_type_coder_tuples()
+    logical_type_reg = schemas.LogicalType._known_logical_types.copy_custom()
+
+    pickler = cloudpickle.CloudPickler(file)
+    # TODO(https://github.com/apache/beam/issues/18500) add file system registry
+    # once implemented
+    pickler.dump({"coder": coder_reg, "logical_type": logical_type_reg})
 
 
 def load_session(file_path):
-  # It is possible to load_session with cloudpickle. However, since references
-  # are saved it should not be necessary. See https://s.apache.org/beam-picklers
-  pass
+  from apache_beam.coders import typecoders
+  from apache_beam.typehints import schemas
+
+  with _pickle_lock, open(file_path, 'rb') as file:
+    registries = cloudpickle.load(file)
+    if type(registries) != dict:
+      raise ValueError(
+          "Faled loading session: expected dict, got {}", type(registries))
+    if "coder" in registries:
+      typecoders.registry.load_custom_type_coder_tuples(registries["coder"])
+    else:
+      _LOGGER.warning('No coder registry found in saved session')
+    if "logical_type" in registries:
+      schemas.LogicalType._known_logical_types.load(registries["logical_type"])
+    else:
+      _LOGGER.warning('No logical type registry found in saved session')

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/apache_beam/internal/cloudpickle_pickler_test.py

@@ -20,6 +20,7 @@
 # pytype: skip-file
 
 import os
+import tempfile
 import threading
 import types
 import unittest
@@ -31,6 +32,7 @@ from apache_beam.internal import code_object_pickler
 from apache_beam.internal import module_test
 from apache_beam.internal.cloudpickle_pickler import dumps
 from apache_beam.internal.cloudpickle_pickler import loads
+from apache_beam.typehints.schemas import LogicalTypeRegistry
 from apache_beam.utils import shared
 
 GLOBAL_DICT_REF = module_test.GLOBAL_DICT
@@ -244,6 +246,24 @@ self.assertEqual(DataClass(datum='abc'), loads(dumps(DataClass(datum='abc'))))
     unpickled_filename = os.path.abspath(unpickled_code.co_filename)
     self.assertEqual(unpickled_filename, original_filename)
 
+  @mock.patch(
+      "apache_beam.coders.typecoders.registry.load_custom_type_coder_tuples")
+  @mock.patch(
+      "apache_beam.typehints.schemas.LogicalType._known_logical_types.load")
+  def test_dump_load_session(self, logicaltype_mock, coder_mock):
+    session_file = 'pickled'
+
+    with tempfile.TemporaryDirectory() as tmp_dirname:
+      pickled_session_file = os.path.join(tmp_dirname, session_file)
+      beam_cloudpickle.dump_session(pickled_session_file)
+      beam_cloudpickle.load_session(pickled_session_file)
+    load_logical_types = logicaltype_mock.call_args.args
+    load_coders = coder_mock.call_args.args
+    self.assertEqual(len(load_logical_types), 1)
+    self.assertEqual(len(load_coders), 1)
+    self.assertTrue(isinstance(load_logical_types[0], LogicalTypeRegistry))
+    self.assertTrue(isinstance(load_coders[0], list))
+
 
 if __name__ == '__main__':
   unittest.main()

{apache_beam-2.70.0rc4 → apache_beam-2.71.0}/apache_beam/internal/pickler.py

@@ -91,6 +91,14 @@ def load_session(file_path):
   return desired_pickle_lib.load_session(file_path)
 
 
+def is_currently_dill():
+  return desired_pickle_lib == dill_pickler
+
+
+def is_currently_cloudpickle():
+  return desired_pickle_lib == cloudpickle_pickler
+
+
 def set_library(selected_library=DEFAULT_PICKLE_LIB):
   """ Sets pickle library that will be used. """
   global desired_pickle_lib
@@ -108,12 +116,11 @@ def set_library(selected_library=DEFAULT_PICKLE_LIB):
         "Pipeline option pickle_library=dill_unsafe is set, but dill is not "
         "installed. Install dill in job submission and runtime environments.")
 
-  is_currently_dill = (desired_pickle_lib == dill_pickler)
   dill_is_requested = (
       selected_library == USE_DILL or selected_library == USE_DILL_UNSAFE)
 
   # If switching to or from dill, update the pickler hook overrides.
-  if is_currently_dill != dill_is_requested:
+  if is_currently_dill() != dill_is_requested:
     dill_pickler.override_pickler_hooks(selected_library == USE_DILL)
 
   if dill_is_requested: