[Python] Fix rebatching into multi-column pyarrow Tables (#41042)

CHANGES.md

@@ -127,6 +127,7 @@
 
 ## Bugfixes
 
+* (Python) Fixed batched DoFns whose input must be rebatched into a `pyarrow.Table` with more than one column, which failed with `ArrowInvalid: Column 1 ... expected length N but got length 0`. This affected, for example, Beam YAML `ReadFromCsv` or `ReadFromJson` followed directly by `WriteToParquet` ([#41042](https://github.com/apache/beam/issues/41042)).
 * (Python) RunInference now measures model memory (the `model_byte_size` metric) when `torch` is not installed. It was always 0, because a failed `ModelManager` import also disabled the `resource` module ([#40854](https://github.com/apache/beam/issues/40854)).
 * (Python/YAML) Fixed `WindowInto` `'d'` (day) duration unit parsing (`24h` instead of `12h`) and fixed `ExtractWindowingInfo` duplicate field name validation ([#40902](https://github.com/apache/beam/issues/40902)).
 * (Python) Fixed `GeminiModelHandler` returning only one result per batch: each prompt is now sent as its own request (up to 32 at a time), so every input gets its own result. A custom `request_fn` must now return one response per input ([#40913](https://github.com/apache/beam/issues/40913)).

sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py

@@ -207,6 +207,56 @@ def test_batch_rebatch_pardos(self):
 
         assert_that(res, equal_to([9, 15, 21]))
 
+  @pytest.mark.uses_pyarrow
+  def test_batch_rebatch_pardos_into_multi_column_batches(self):
+    # Rebatching passes the producer's exploded elements to the consumer's
+    # batch converter, which may read them more than once: the pyarrow Table
+    # converter reads them once per column.
+    import pyarrow as pa
+
+    from apache_beam.typehints.row_type import RowTypeConstraint
+
+    class ListBatchDoFn(beam.DoFn):
+      def process_batch(self, batch, *unused_args, **unused_kwargs):
+        assert isinstance(batch, list)
+        yield batch
+
+      def infer_output_type(self, input_type):
+        return input_type
+
+      def get_input_batch_type(self, input_element_type):
+        return list[input_element_type]
+
+      def get_output_batch_type(self, input_element_type):
+        return list[input_element_type]
+
+    class TableBatchDoFn(beam.DoFn):
+      def process_batch(self, batch, *unused_args, **unused_kwargs):
+        assert isinstance(batch, pa.Table)
+        yield batch
+
+      def infer_output_type(self, input_type):
+        return input_type
+
+      def get_input_batch_type(self, input_element_type):
+        return pa.Table
+
+      def get_output_batch_type(self, input_element_type):
+        return pa.Table
+
+    row_type = RowTypeConstraint.from_fields([('id', int), ('name', str)])
+    with self.create_pipeline() as p:
+      res = (
+          p
+          | beam.Create([
+              beam.Row(id=1, name='Ada'), beam.Row(id=2, name='Grace')
+          ]).with_output_types(row_type)
+          | beam.ParDo(ListBatchDoFn())
+          | beam.ParDo(TableBatchDoFn())
+          | beam.Map(lambda row: (row.id, row.name)))
+
+      assert_that(res, equal_to([(1, 'Ada'), (2, 'Grace')]))
+
   def test_batch_pardo_fusion_break(self):
     class NormalizeDoFn(beam.DoFn):
       @no_type_check

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

@@ -383,11 +383,14 @@ def receive_batch(self, windowed_batch):
             "This is very inefficient, consider re-structuring your pipeline "
             "or adding a DoFn to directly convert between these types.",
             InefficientExecutionWarning)
+        # produce_batch takes a sequence, and some converters read it more
+        # than once (the pyarrow Table converter reads it once per column).
         _cast_to_operation(consumer).process_batch(
             windowed_batch.with_values(
                 consumer_batch_converter.produce_batch(
-                    self.producer_batch_converter.explode_batch(
-                        windowed_batch.values))))
+                    list(
+                        self.producer_batch_converter.explode_batch(
+                            windowed_batch.values)))))
 
     self.update_counters_batch(windowed_batch)