diff --git a/CHANGES.md b/CHANGES.md index 61b1547e025f..33c2ac650fc8 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -63,6 +63,7 @@ ## Bugfixes +* (Python/YAML) Fixed SQL-backed Combine generating invalid syntax for Calcite expressions and SQL shorthand aggregations ([#41288](https://github.com/apache/beam/issues/41288)). * (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..4e92ca205514 100644 --- a/sdks/python/apache_beam/yaml/yaml_combine.py +++ b/sdks/python/apache_beam/yaml/yaml_combine.py @@ -62,15 +62,22 @@ def normalize_combine(spec): type: fn_type """ from apache_beam.yaml.yaml_transform import SafeLineLoader - if spec['type'] == 'Combine': + if spec['type'] == 'Combine' or spec['type'].startswith('Combine-'): config = spec.get('config') if isinstance(config.get('group_by'), str): config['group_by'] = [config['group_by']] + is_sql = ( + config.get('language') in ('sql', 'calcite') or + spec['type'] in ('Combine-sql', 'Combine-calcite')) + def normalize_agg(dest, agg): if isinstance(agg, str): agg = {'fn': agg} - if 'value' not in agg and config.get('language') != 'sql': + fn_str = ( + agg.get('fn') if isinstance(agg.get('fn'), str) else agg.get( + 'fn', {}).get('type', '')) + if 'value' not in agg and not (is_sql and '(' in fn_str): agg['value'] = dest if isinstance(agg['fn'], str): agg['fn'] = {'type': agg['fn']} @@ -189,13 +196,18 @@ def _SqlCombineTransform( raise ValueError(f'Unknown grouping columns: {list(unknown_keys)}') def combine_col(dest, fn_spec): + fn_type = fn_spec['fn']['type'] if 'value' in fn_spec or 'config' in fn_spec['fn']: - expr = '%s(%s)' % ( - fn_spec['fn']['type'], - ', '.join([fn_spec['value']] + - list(fn_spec['fn'].get('config', {}).values()))) + if '(' in fn_type: + expr = fn_type + else: + args = [fn_spec.get('value', dest) + ] + [str(v) for v in fn_spec['fn'].get('config', {}).values()] + expr = '%s(%s)' % (fn_type, ', '.join(args)) + elif '(' not in fn_type: + expr = f'{fn_type}({dest})' else: - expr = fn_spec['fn']['type'] + expr = fn_type return f'{expr} as {dest}' return pcoll | sql_transform_constructor( diff --git a/sdks/python/apache_beam/yaml/yaml_combine_test.py b/sdks/python/apache_beam/yaml/yaml_combine_test.py index caf3de10078b..f0b0bd7e3267 100644 --- a/sdks/python/apache_beam/yaml/yaml_combine_test.py +++ b/sdks/python/apache_beam/yaml/yaml_combine_test.py @@ -21,6 +21,8 @@ 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_combine import normalize_combine from apache_beam.yaml.yaml_transform import YamlTransform DATA = [ @@ -186,6 +188,82 @@ def test_config(self): beam.Row(b=2, biggest=[104]), ])) + def test_sql_combine_calcite_expression(self): + queries = [] + pcoll = beam.Pipeline() | beam.Create(DATA) + spec = { + 'type': 'Combine', + 'config': { + 'language': 'calcite', + 'group_by': ['a'], + 'combine': { + 'total_c': 'sum(c)', 'num_records': 'count(*)' + } + } + } + normalized = normalize_combine(spec) + _ = pcoll | _SqlCombineTransform( + lambda q: queries.append(q) or beam.Map(lambda x: x), + **normalized['config']) + self.assertEqual( + queries[-1], + 'SELECT a, sum(c) as total_c, count(*) as num_records ' + 'FROM PCOLLECTION GROUP BY a') + + def test_sql_combine_sql_shorthand(self): + queries = [] + pcoll = beam.Pipeline() | beam.Create(DATA) + spec = { + 'type': 'Combine', + 'config': { + 'language': 'sql', 'group_by': ['a'], 'combine': { + 'c': 'sum' + } + } + } + normalized = normalize_combine(spec) + _ = pcoll | _SqlCombineTransform( + lambda q: queries.append(q) or beam.Map(lambda x: x), + **normalized['config']) + self.assertEqual( + queries[-1], 'SELECT a, sum(c) as c FROM PCOLLECTION GROUP BY a') + + def test_sql_combine_calcite_shorthand(self): + queries = [] + pcoll = beam.Pipeline() | beam.Create(DATA) + spec = { + 'type': 'Combine', + 'config': { + 'language': 'calcite', 'group_by': ['a'], 'combine': { + 'c': 'sum' + } + } + } + normalized = normalize_combine(spec) + _ = pcoll | _SqlCombineTransform( + lambda q: queries.append(q) or beam.Map(lambda x: x), + **normalized['config']) + self.assertEqual( + queries[-1], 'SELECT a, sum(c) as c FROM PCOLLECTION GROUP BY a') + + def test_sql_combine_type_with_language_suffix(self): + queries = [] + pcoll = beam.Pipeline() | beam.Create(DATA) + spec = { + 'type': 'Combine-calcite', + 'config': { + 'group_by': ['a'], 'combine': { + 'total_c': 'sum(c)' + } + } + } + normalized = normalize_combine(spec) + _ = pcoll | _SqlCombineTransform( + lambda q: queries.append(q) or beam.Map(lambda x: x), + **normalized['config']) + self.assertEqual( + queries[-1], 'SELECT a, sum(c) as total_c FROM PCOLLECTION GROUP BY a') + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO)