diff --git a/CHANGES.md b/CHANGES.md index 1eb1229261ad..ab7e2c03b542 100644 --- a/CHANGES.md +++ b/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)). diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py index 255008ab13f8..8566accd0ea4 100644 --- a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py +++ b/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 diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py index bfa1c0f3d47c..82a54cbcd3b2 100644 --- a/sdks/python/apache_beam/runners/worker/operations.py +++ b/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)