[YAML] Fix SQL-backed Combine query generation for Calcite expressions and SQL shorthand (#41288)

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)).
 

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(

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)