Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)).

Expand Down
42 changes: 28 additions & 14 deletions sdks/python/apache_beam/yaml/yaml_combine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -101,13 +103,15 @@ class PyJsYamlCombine(beam.PTransform):
language: The language used to define (and execute) the
custom callables in `combine`. Defaults to generic.
"""

# pylint: disable=dangerous-default-value
def __init__(
self,
group_by: Iterable[str],
combine: Mapping[str, Mapping[str, Any]],
group_by: Optional[Iterable[str]] = (),
combine: Mapping[str, Mapping[str, Any]] = {},
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):
Expand Down Expand Up @@ -178,9 +182,15 @@ def extract_return_type(expr):
', '.join('`%s`' % k for k in BUILTIN_COMBINE_FNS.keys()))


# pylint: disable=dangerous-default-value
@beam.ptransform.ptransform_fn
def _SqlCombineTransform(
pcoll, sql_transform_constructor, group_by, combine, language=None):
pcoll,
sql_transform_constructor,
group_by: Optional[Iterable[str]] = (),
combine: Mapping[str, Any] = {},
language: Optional[str] = None):
group_by = list(group_by or [])
all_fields = [
x for x, _ in named_fields_from_element_type(pcoll.element_type)
]
Expand All @@ -198,14 +208,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():
Expand Down
56 changes: 56 additions & 0 deletions sdks/python/apache_beam/yaml/yaml_combine_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -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:
Expand Down
Loading