apache-beam 2.76.0rc4__tar.gz → 2.77.0rc2__tar.gz

{apache_beam-2.76.0rc4 → apache_beam-2.77.0rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-beam
-Version: 2.76.0rc4
+Version: 2.77.0rc2
 Summary: Apache Beam SDK for Python
 Home-page: https://beam.apache.org
 Download-URL: https://pypi.python.org/pypi/apache-beam
@@ -20,7 +20,7 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: cryptography<49.0.0,>=39.0.0
+Requires-Dist: cryptography<51.0.0,>=39.0.0
 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.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*,<2,>=1.33.1
@@ -85,6 +85,7 @@ Requires-Dist: sqlalchemy-pytds>=1.0.2; extra == "test"
 Requires-Dist: pg8000>=1.31.5; extra == "test"
 Requires-Dist: PyMySQL>=1.1.0; extra == "test"
 Requires-Dist: oracledb>=3.1.1; extra == "test"
+Requires-Dist: deltalake>=0.15.0; extra == "test"
 Provides-Extra: gcp
 Requires-Dist: cachetools<7,>=3.1.0; extra == "gcp"
 Requires-Dist: google-api-core<3,>=2.0.0; extra == "gcp"
@@ -100,7 +101,7 @@ Requires-Dist: google-cloud-dataflow-client<0.14.0,>=0.13.0; extra == "gcp"
 Requires-Dist: google-cloud-bigquery<4,>=2.0.0; extra == "gcp"
 Requires-Dist: google-cloud-bigquery-storage<3,>=2.6.3; extra == "gcp"
 Requires-Dist: google-cloud-core<3,>=2.0.0; extra == "gcp"
-Requires-Dist: google-cloud-bigtable<3,>=2.42.0; extra == "gcp"
+Requires-Dist: google-cloud-bigtable<3,>=2.44.0; extra == "gcp"
 Requires-Dist: google-cloud-build<4,>=3.35.0; extra == "gcp"
 Requires-Dist: google-cloud-spanner<4,>=3.0.0; extra == "gcp"
 Requires-Dist: google-cloud-dlp<4,>=3.0.0; extra == "gcp"
@@ -118,6 +119,8 @@ Requires-Dist: PyMySQL>=1.1.0; extra == "gcp"
 Requires-Dist: keyrings.google-artifactregistry-auth; extra == "gcp"
 Requires-Dist: orjson<4,>=3.9.7; extra == "gcp"
 Requires-Dist: regex>=2020.6.8; extra == "gcp"
+Provides-Extra: gcp-test
+Requires-Dist: google-cloud-firestore<3,>=2.0.0; extra == "gcp-test"
 Provides-Extra: interactive
 Requires-Dist: facets-overview<2,>=1.1.0; extra == "interactive"
 Requires-Dist: google-cloud-dataproc<6,>=5.0.0; extra == "interactive"

{apache_beam-2.76.0rc4 → apache_beam-2.77.0rc2}/apache_beam/dataframe/frames.py

@@ -1091,27 +1091,81 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
       reindexed = self.reorder_levels(
           level + [i for i in range(self.index.nlevels) if i not in level])
 
-    def xs_partitioned(frame, key):
-      if not len(key):
-        # key is not in this partition, return empty dataframe
-        result = frame.iloc[:0]
-        if key_size < frame.index.nlevels:
+    if key_size < reindexed.index.nlevels:
+
+      def xs_partitioned(frame, key):
+        if not len(key):
+          # key is not in this partition, return empty dataframe/series
+          result = frame.iloc[:0]
           return result.droplevel(list(range(key_size)))
-        else:
-          return result
+        return frame.xs(key.item(), **kwargs)
 
-      # key should be in this partition, call xs. Will raise KeyError if not
-      # present.
-      return frame.xs(key.item())
+      return frame_base.DeferredFrame.wrap(
+          expressions.ComputedExpression(
+              'xs',
+              xs_partitioned, [reindexed._expr, key_expr],
+              requires_partition_by=partitionings.Index(list(range(key_size))),
+              preserves_partition_by=partitionings.Singleton()))
+    else:
+      # When all index levels are matched (key_size >= nlevels), pandas .xs()
+      # return type is data-dependent:
+      #   - Single match: reduces dimensionality (DataFrame -> Series, Series -> scalar)
+      #   - Duplicate matches: preserves container type (DataFrame -> DataFrame, Series -> Series)
+      # Because proxy schemas are 0-row templates evaluated at graph construction time
+      # without knowledge of dataset contents or key frequencies, the proxy always assumes
+      # a single match (dimensionality-reduced type). At runtime, the Singleton unwrap stage
+      # correctly produces whichever type pandas returns. Tests with multi-matching keys
+      # therefore specify check_proxy=False.
+      def xs_partitioned_wrapped(frame, key):
+        if not len(key):
+          return pd.Series([], dtype=object)
+        k = key.item()
+        try:
+          res = frame.xs(k, **kwargs)
+          return pd.Series([res], dtype=object)
+        except KeyError:
+          return pd.Series([], dtype=object)
+
+      intermediate = expressions.ComputedExpression(
+          'xs_partitioned_wrapped',
+          xs_partitioned_wrapped, [reindexed._expr, key_expr],
+          proxy=pd.Series([], dtype=object),
+          requires_partition_by=partitionings.Index(list(range(key_size))),
+          preserves_partition_by=partitionings.Singleton())
 
-    return frame_base.DeferredFrame.wrap(
-        expressions.ComputedExpression(
-            'xs',
-            xs_partitioned,
-            [reindexed._expr, key_expr],
-            requires_partition_by=partitionings.Index(list(range(key_size))),
-            # Drops index levels, so partitioning is not preserved
-            preserves_partition_by=partitionings.Singleton()))
+      proxy_frame = reindexed._expr.proxy().iloc[:0]
+      if not proxy_frame.index.is_unique:
+        proxy_frame.index = proxy_frame.index.drop_duplicates()
+      k_val = key_series.iloc[0]
+      dummy_index = (
+          pd.MultiIndex.from_tuples([k_val], names=proxy_frame.index.names) if
+          isinstance(k_val, tuple) else pd.Index([k_val],
+                                                 name=proxy_frame.index.name))
+
+      if isinstance(proxy_frame, pd.DataFrame):
+        dummy_obj = proxy_frame.reindex(dummy_index)
+        xs_proxy = dummy_obj.xs(k_val, **kwargs)
+        if isinstance(xs_proxy, (pd.DataFrame, pd.Series)):
+          xs_proxy = xs_proxy.iloc[:0]
+      else:
+        try:
+          xs_proxy = proxy_frame.dtype.type()
+        except TypeError:
+          xs_proxy = proxy_frame.reindex(dummy_index).iloc[0]
+
+      def unwrap_xs(ser):
+        if ser.empty:
+          raise KeyError(k_val)
+        return ser.iloc[0]
+
+      with expressions.allow_non_parallel_operations(True):
+        return frame_base.DeferredFrame.wrap(
+            expressions.ComputedExpression(
+                'xs',
+                unwrap_xs, [intermediate],
+                proxy=xs_proxy,
+                requires_partition_by=partitionings.Singleton(),
+                preserves_partition_by=partitionings.Singleton()))
 
   @property
   def dtype(self):

{apache_beam-2.76.0rc4 → apache_beam-2.77.0rc2}/apache_beam/dataframe/frames_test.py

@@ -331,6 +331,19 @@ class DeferredFrameTest(_AbstractFrameTest):
         lambda df: df.num_legs.xs(('bird', 'walks'), level=[0, 'locomotion']),
         df)
 
+    # Test cases reported in BEAM-28559
+    df_single_index = df.reset_index().set_index('class')
+    self._run_test(
+        lambda df: df.num_legs.xs('mammal'), df_single_index, check_proxy=False)
+    self._run_test(lambda df: df.num_legs.xs('bird'), df_single_index)
+
+    # Categorical Series single match
+    s_cat = pd.Series(
+        pd.Categorical(['a', 'b', 'c']),
+        index=['r1', 'r2', 'r3'],
+        name='cat_col')
+    self._run_test(lambda s: s.xs('r1'), s_cat, check_proxy=False)
+
   def test_dataframe_xs(self):
     # Test cases reported in BEAM-13421
     df = pd.DataFrame(
@@ -342,10 +355,67 @@ class DeferredFrameTest(_AbstractFrameTest):
         ]),
         columns=['provider', 'time', 'value'])
 
-    self._run_test(lambda df: df.xs('state'), df.set_index(['provider']))
+    self._run_test(
+        lambda df: df.xs('state'),
+        df.set_index(['provider']),
+        check_proxy=False)
     self._run_test(
         lambda df: df.xs('state'), df.set_index(['provider', 'time']))
 
+    # Test cases reported in BEAM-28559
+    self._run_test(lambda df: df.xs('county'), df.set_index(['provider']))
+    self._run_test(
+        lambda df: df.xs(('state', 'day1')),
+        df.set_index(['provider', 'time']),
+        check_proxy=False)
+
+    df_unique = pd.DataFrame(
+        np.array([
+            ['state', 'day1', 12],
+            ['state', 'day2', 14],
+            ['county', 'day1', 9],
+        ]),
+        columns=['provider', 'time', 'value'])
+    self._run_test(
+        lambda df: df.xs(('state', 'day2')),
+        df_unique.set_index(['provider', 'time']))
+
+    # Categorical and extension dtype tests
+    df_cat = pd.DataFrame({
+        'cat': pd.Categorical(['a', 'b', 'c']), 'val': [1, 2, 3]
+    },
+                          index=['r1', 'r2', 'r3'])
+    self._run_test(lambda df: df.xs('r1'), df_cat)
+
+    df_dt_tz = pd.DataFrame({
+        'dt': pd.Series([
+            pd.Timestamp('2023-01-01', tz='UTC'),
+            pd.Timestamp('2023-01-02', tz='UTC')
+        ],
+                        dtype='datetime64[ns, UTC]'),
+        'val': [1, 2]
+    },
+                            index=['r1', 'r2'])
+    self._run_test(lambda df: df.xs('r1'), df_dt_tz)
+
+    df_null_int = pd.DataFrame({'num': pd.Series([1, 2, None], dtype='Int64')},
+                               index=['r1', 'r2', 'r3'])
+    self._run_test(lambda df: df.xs('r1'), df_null_int)
+
+  def test_dataframe_xs_non_empty_duplicate_proxy(self):
+    df_dups = pd.DataFrame({'a': [1, 2]}, index=['x', 'x'])
+    p = beam.Pipeline()
+    deferred = to_dataframe(p | beam.Create([{'a': 1}]), proxy=df_dups)
+    res = deferred.xs('x')
+    self.assertIsInstance(res, frames.DeferredSeries)
+    self.assertTrue(res._expr.proxy().empty)
+
+    s_dups = pd.Series(pd.Categorical(['a', 'b']), index=['x', 'x'], name='s')
+    deferred_s = to_dataframe(
+        p | 'CreateSeries' >> beam.Create(['a']), proxy=s_dups)
+    res_s = deferred_s.xs('x')
+    self.assertIsInstance(res_s, frame_base.DeferredBase)
+
   def test_set_column(self):
     def new_column(df):
       df['NewCol'] = df['Speed']

{apache_beam-2.76.0rc4 → apache_beam-2.77.0rc2}/apache_beam/examples/inference/pytorch_image_object_detection.py

@@ -431,7 +431,8 @@ def run(
       model_params={},
       state_dict_path=known_args.model_state_dict_path,
       device=device,
-      inference_batch_size=batch_size,
+      min_batch_size=batch_size,
+      max_batch_size=batch_size,
       inference_fn=_torchvision_detection_inference_fn,
   )
 

{apache_beam-2.76.0rc4 → apache_beam-2.77.0rc2}/apache_beam/examples/inference/runinference_metrics/setup.py

@@ -29,7 +29,7 @@ import setuptools
 from setuptools import find_packages
 
 REQUIREMENTS = [
-    "apache-beam[gcp]==2.41.0", "transformers==4.38.0", "torch==1.13.1"
+    "apache-beam[gcp]==2.41.0", "transformers==5.10.1", "torch==1.13.1"
 ]
 
 setuptools.setup(

{apache_beam-2.76.0rc4 → apache_beam-2.77.0rc2}/apache_beam/io/external/xlang_jdbcio_it_test.py

@@ -18,8 +18,10 @@
 # pytype: skip-file
 
 import datetime
+import json
 import logging
 import os
+import sys
 import time
 import typing
 import unittest
@@ -46,6 +48,13 @@ except ImportError:
   sqlalchemy = None
 # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
 
+# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
+try:
+  from google.cloud import secretmanager
+except ImportError:
+  secretmanager = None  # type: ignore[assignment]
+# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
+
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
 try:
   from testcontainers.mysql import MySqlContainer
@@ -138,6 +147,42 @@ class CrossLanguageJdbcIOTest(unittest.TestCase):
     cls.engines = {}
     cls.jdbc_configs = {}
 
+    cls.secret_manager_available = False
+    if secretmanager is not None:
+      try:
+        cls.project_id = os.environ.get(
+            'GOOGLE_CLOUD_PROJECT', 'apache-beam-testing')
+        cls.secret_client = secretmanager.SecretManagerServiceClient()
+        py_version = f'_py{sys.version_info.major}{sys.version_info.minor}'
+        secret_postfix = (
+            datetime.datetime.now().strftime('%m%d_%H%M%S') + py_version)
+        cls.secret_id = 'xlang_jdbc_test_secret_' + secret_postfix
+        cls.project_path = f'projects/{cls.project_id}'
+        cls.secret_path = f'{cls.project_path}/secrets/{cls.secret_id}'
+        try:
+          cls.secret_client.get_secret(request={'name': cls.secret_path})
+        except Exception:
+          cls.secret_client.create_secret(
+              request={
+                  'parent': cls.project_path,
+                  'secret_id': cls.secret_id,
+                  'secret': {
+                      'replication': {
+                          'automatic': {}
+                      }
+                  }
+              })
+        cls.secret_client.add_secret_version(
+            request={
+                'parent': cls.secret_path, 'payload': {
+                    'data': b'test'
+                }
+            })
+        cls.secret_manager_available = True
+      except Exception as e:
+        logging.warning("Could not set up GCP Secret Manager: %s", e)
+        cls.secret_manager_available = False
+
     for db_type, db_data in cls.DB_CONTAINER_CLASSPATH_STRING.items():
       container = cls.start_container(db_data.container_fn)
       cls.containers[db_type] = container
@@ -162,6 +207,13 @@ class CrossLanguageJdbcIOTest(unittest.TestCase):
 
   @classmethod
   def tearDownClass(cls):
+    if getattr(cls, 'secret_manager_available',
+               False) and secretmanager is not None:
+      try:
+        cls.secret_client.delete_secret(request={'name': cls.secret_path})
+      except Exception:  # pylint: disable=broad-except
+        logging.warning("Could not delete GCP secret: %s", cls.secret_path)
+
     for db_type, container in cls.containers.items():
       if container:
         # Sometimes stopping the container raises ReadTimeout. We can ignore it
@@ -454,6 +506,67 @@ class CrossLanguageJdbcIOTest(unittest.TestCase):
 
       assert_that(result, equal_to(expected_filtered_rows))
 
+  @parameterized.expand(['postgres', 'mysql'])
+  def test_xlang_jdbc_with_secret_manager(self, database):
+    if not getattr(self, 'secret_manager_available', False):
+      self.skipTest(
+          "GCP Secret Manager is not available or credentials not configured.")
+
+    if self.containers[database] is None:
+      self.skipTest(f"{database} container could not be initialized")
+
+    table_name = f"jdbc_secret_manager_test_{database}"
+
+    with self.engines[database].begin() as connection:
+      connection.execute(
+          sqlalchemy.text(
+              f"CREATE TABLE IF NOT EXISTS {table_name}" +
+              "(id INTEGER, name VARCHAR(50), value DOUBLE PRECISION)"))
+
+    test_rows = [
+        SimpleRow(1, "Item1", 10.5),
+        SimpleRow(2, "Item2", 20.75),
+        SimpleRow(3, "Item3", 30.25),
+        SimpleRow(4, "Item4", 40.0),
+        SimpleRow(-5, "Item5", 50.5),
+    ]
+
+    config = self.jdbc_configs[database]
+    secret_password_spec = json.dumps({
+        'name': self.secret_id, 'project': self.project_id
+    })
+
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      _ = (
+          p
+          | beam.Create(test_rows).with_output_types(SimpleRow)
+          | 'Write to jdbc with secret manager' >> WriteToJdbc(
+              table_name=table_name,
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=secret_password_spec,
+              secret_manager='googlecloudsecretmanager',
+              classpath=config['classpath'],
+          ))
+
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      result = (
+          p
+          | 'Read from jdbc with secret manager' >> ReadFromJdbc(
+              table_name=table_name,
+              driver_class_name=config['driver_class_name'],
+              jdbc_url=config['jdbc_url'],
+              username=config['username'],
+              password=secret_password_spec,
+              secret_manager='googlecloudsecretmanager',
+              classpath=config['classpath'],
+              schema=SimpleRow))
+
+      assert_that(result, equal_to(test_rows))
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)

{apache_beam-2.76.0rc4 → apache_beam-2.77.0rc2}/apache_beam/io/external/xlang_jmsio_it_test.py

@@ -92,7 +92,7 @@ class _BaseJmsIOTest(unittest.TestCase):
     subscriber_result = {}
 
     def publish():
-      self.produce(source_queue, remaining_records)
+      self.produce(source_queue, NUM_RECORDS)
 
     stop_event = threading.Event()
 
@@ -110,12 +110,6 @@ class _BaseJmsIOTest(unittest.TestCase):
           break
       _LOGGER.info('received %s messages', len(received_messages))
 
-    # TODO(https://github.com/apache/beam/issues/39446): Clean up
-    # pre-publishing Prism runner issue resolved
-    initial_records = 10
-    remaining_records = NUM_RECORDS - initial_records
-    self.produce(source_queue, initial_records)
-
     publisher = threading.Thread(target=publish, daemon=True)
     subscriber = threading.Thread(target=subscribe, daemon=True)
 
@@ -147,7 +141,7 @@ class _BaseJmsIOTest(unittest.TestCase):
       result = p.run()
       subscriber.start()
       try:
-        subscriber.join(timeout=20)  # 1.5 min
+        subscriber.join(timeout=20)
       finally:
         stop_event.set()
         publisher.join()
@@ -160,8 +154,7 @@ class _BaseJmsIOTest(unittest.TestCase):
 
     received = subscriber_result.get('received', [])
     self.assertEqual(len(received), NUM_RECORDS)
-    # there are identical records
-    self.assertEqual(len(set(received)), NUM_RECORDS - initial_records)
+    self.assertEqual(len(set(received)), NUM_RECORDS)
 
 
 class ActiveMQJmsIOTest(_BaseJmsIOTest):