diff --git a/CHANGES.md b/CHANGES.md index 1eb1229261ad..ec998f6c1c16 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -63,6 +63,7 @@ ## Bugfixes +* (Python/YAML) Fixed SQL-backed `Combine` generating an invalid `GROUP BY ` clause for global aggregations when `group_by` is empty or omitted ([#41047](https://github.com/apache/beam/issues/41047)). * (Java) Improved the error message when decoding with a ProtoCoder for the raw `com.google.protobuf.Message` type (for example inferred from a generic `T extends Message`): it now explains how to fix it instead of surfacing a `NoSuchMethodException` ([#19713](https://github.com/apache/beam/issues/19713)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). diff --git a/sdks/python/apache_beam/yaml/yaml_combine.py b/sdks/python/apache_beam/yaml/yaml_combine.py index d32b5272106e..e5452f3ec9f9 100644 --- a/sdks/python/apache_beam/yaml/yaml_combine.py +++ b/sdks/python/apache_beam/yaml/yaml_combine.py @@ -64,7 +64,9 @@ def normalize_combine(spec): from apache_beam.yaml.yaml_transform import SafeLineLoader if spec['type'] == 'Combine': config = spec.get('config') - if isinstance(config.get('group_by'), str): + if not config.get('group_by'): + config['group_by'] = [] + elif isinstance(config['group_by'], str): config['group_by'] = [config['group_by']] def normalize_agg(dest, agg): @@ -103,11 +105,11 @@ class PyJsYamlCombine(beam.PTransform): """ def __init__( self, - group_by: Iterable[str], - combine: Mapping[str, Mapping[str, Any]], + group_by: Iterable[str] = (), + combine: Optional[Mapping[str, Mapping[str, Any]]] = None, language: Optional[str] = None): - self._group_by = group_by - self._combine = combine + self._group_by = tuple(group_by or ()) + self._combine = combine or {} self._language = language def expand(self, pcoll): @@ -180,7 +182,12 @@ def extract_return_type(expr): @beam.ptransform.ptransform_fn def _SqlCombineTransform( - pcoll, sql_transform_constructor, group_by, combine, language=None): + pcoll, + sql_transform_constructor, + group_by: Iterable[str] = (), + combine: Optional[Mapping[str, Any]] = None, + language: Optional[str] = None): + group_by = list(group_by or []) all_fields = [ x for x, _ in named_fields_from_element_type(pcoll.element_type) ] @@ -198,14 +205,18 @@ def combine_col(dest, fn_spec): expr = fn_spec['fn']['type'] return f'{expr} as {dest}' - return pcoll | sql_transform_constructor( - 'SELECT %s FROM PCOLLECTION GROUP BY %s' % ( - ', '.join( - list(group_by) + - [combine_col(dest, fn_spec) - for dest, fn_spec in combine.items()]), - ', '.join(group_by), - )) + select_cols = group_by + [ + combine_col(dest, fn_spec) for dest, fn_spec in (combine or {}).items() + ] + if group_by: + query = 'SELECT %s FROM PCOLLECTION GROUP BY %s' % ( + ', '.join(select_cols), + ', '.join(group_by), + ) + else: + query = 'SELECT %s FROM PCOLLECTION' % ', '.join(select_cols) + + return pcoll | sql_transform_constructor(query) def create_combine_providers(): diff --git a/sdks/python/apache_beam/yaml/yaml_combine_test.py b/sdks/python/apache_beam/yaml/yaml_combine_test.py index caf3de10078b..a9deff6c9a95 100644 --- a/sdks/python/apache_beam/yaml/yaml_combine_test.py +++ b/sdks/python/apache_beam/yaml/yaml_combine_test.py @@ -21,6 +21,7 @@ import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to +from apache_beam.yaml.yaml_combine import _SqlCombineTransform from apache_beam.yaml.yaml_transform import YamlTransform DATA = [ @@ -90,6 +91,61 @@ def test_no_keys(self): beam.Row(c=410), ])) + def test_no_keys_omitted(self): + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: + elements = p | beam.Create(DATA) + result = elements | YamlTransform( + ''' + type: Combine + config: + combine: + c: sum + ''') + assert_that( + result | beam.Map(lambda x: beam.Row(**x._asdict())), + equal_to([ + beam.Row(c=410), + ])) + + def test_sql_combine_query_generation(self): + queries = [] + pcoll = beam.Pipeline() | beam.Create(DATA) + + # With grouping keys: GROUP BY clause must be generated. + _ = pcoll | _SqlCombineTransform( + lambda q: queries.append(q) or beam.Map(lambda x: x), + group_by=['a'], + combine={'total_c': { + 'fn': { + 'type': 'sum' + }, 'value': 'c' + }}) + self.assertEqual( + queries[-1], 'SELECT a, sum(c) as total_c FROM PCOLLECTION GROUP BY a') + + # Without grouping keys (group_by=[]): GROUP BY clause must be omitted. + _ = pcoll | _SqlCombineTransform( + lambda q: queries.append(q) or beam.Map(lambda x: x), + group_by=[], + combine={'total_c': { + 'fn': { + 'type': 'sum' + }, 'value': 'c' + }}) + self.assertEqual(queries[-1], 'SELECT sum(c) as total_c FROM PCOLLECTION') + + # With omitted grouping keys (group_by=None): GROUP BY clause must be omitted. + _ = pcoll | _SqlCombineTransform( + lambda q: queries.append(q) or beam.Map(lambda x: x), + group_by=None, + combine={'total_c': { + 'fn': { + 'type': 'sum' + }, 'value': 'c' + }}) + self.assertEqual(queries[-1], 'SELECT sum(c) as total_c FROM PCOLLECTION') + def test_multiple_combines(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: