Fix bug when execute DataframeTransform a dictionary of Pcoll

sdks/python/apache_beam/dataframe/transforms.py

@@ -114,8 +114,9 @@ def expand(self, input_pcolls):
         for tag in input_dict
     }
     input_frames = {
-        k: convert.to_dataframe(pc, proxies[k])
-        for k, pc in input_dict.items()
+        k: convert.to_dataframe(pc, proxies[k], str(k))
+        for k,
+        pc in input_dict.items()
     }  # type: Dict[Any, DeferredFrame] # noqa: F821
 
     # Apply the function.

sdks/python/apache_beam/dataframe/transforms_test.py

@@ -318,6 +318,45 @@ def check(actual):
           lambda x: {'res': 3 * x}, proxy, yield_elements='pandas')
       assert_that(res['res'], equal_to_series(three_series), 'CheckDictOut')
 
+  def test_multiple_dataframe_transforms(self):
+    # Define test data
+    data1 = [
+        beam.Row(id=1, name="abc"),
+        beam.Row(id=2, name="def"),
+        beam.Row(id=3, name="ghi")
+    ]
+    data2 = [
+        beam.Row(addr="addr1"), beam.Row(addr="addr2"), beam.Row(addr="addr3")
+    ]
+
+    # Create a TestPipeline
+    with beam.Pipeline() as p:
+      # Create PCollections for testing
+      pcol1 = p | "Create1" >> beam.Create(data1)
+      pcol2 = p | "Create2" >> beam.Create(data2)
+
+      # Apply the DataframeTransform to the PCollections
+      pcol = ({
+          "a": pcol1, "b": pcol2
+      }
+              | "TransformedDF" >> transforms.DataframeTransform(
+                  lambda a, b: a.assign(addr="addr-common")))
+
+      # Assert the expected output
+      expected_output = [
+          {
+              "id": 1, "name": "abc", "addr": "addr-common"
+          },
+          {
+              "id": 2, "name": "def", "addr": "addr-common"
+          },
+          {
+              "id": 3, "name": "ghi", "addr": "addr-common"
+          },
+      ]
+      assert_that(pcol | "Map" >> beam.Map(lambda row: row.asdict())) \
+        .equal_to(expected_output)
+
   def test_cat(self):
     # verify that cat works with a List[Series] since this is
     # missing from doctests

sdks/python/apache_beam/transforms/core.py

@@ -1504,8 +1504,8 @@ def _check_fn_use_yield_and_return(fn):
     has_yield = False
     has_return = False
     return_none_warning = (
-        "No iterator is returned by the process method in %s.",
-        fn.__self__.__class__)
+        f"No iterator is returned by the "
+        f"process method in {fn.__self__.__class__}.")
     for line in source_code.split("\n"):
       lstripped_line = line.lstrip()
       if lstripped_line.startswith("yield ") or lstripped_line.startswith(