[Python/YAML] Accept >=, <=, == and != in generic Filter expressions …

CHANGES.md

@@ -130,6 +130,7 @@
 
 ## Bugfixes
 
+* (Python/YAML) Beam YAML `generic` filters now accept `>=`, `<=`, `==` and `!=`; previously only `<` and `>` were recognized ([#41043](https://github.com/apache/beam/issues/41043)).
 * (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)).

sdks/python/apache_beam/yaml/yaml_mapping.py

@@ -125,7 +125,9 @@ def is_atomic(expr: str):
     return
 
   if allow_cmp:
-    maybe_cmp = re.fullmatch('(.*)([<>=!]+)(.*)', expr)
+    # The left side is non-greedy so that two-character operators such as
+    # '>=' stay whole; a greedy match would split 'a >= 1' into 'a >', '='.
+    maybe_cmp = re.fullmatch('(.*?)([<>=!]+)(.*)', expr)
     if maybe_cmp:
       left, cmp, right = maybe_cmp.groups()
       if (is_atomic(left.strip()) and is_atomic(right.strip()) and

sdks/python/apache_beam/yaml/yaml_mapping_test.py

@@ -571,6 +571,59 @@ def test_extract_windowing_info_iterable(self):
                   window_type='GlobalWindow'),
           ]))
 
+  def test_filter_generic(self):
+    with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+        pickle_library='cloudpickle')) as p:
+      elements = p | beam.Create(DATA)
+      result = elements | YamlTransform(
+          '''
+          type: Filter
+          config:
+              language: generic
+              keep: "rank >= 1"
+          ''')
+      assert_that(
+          result
+          | beam.Map(lambda named_tuple: beam.Row(**named_tuple._asdict())),
+          equal_to([
+              beam.Row(label='37a', conductor=37, rank=1),
+              beam.Row(label='389a', conductor=389, rank=2),
+          ]))
+
+
+class ValidateGenericExpressionTest(unittest.TestCase):
+  FIELDS = ['age', 'name']
+
+  def test_accepts_every_comparison_operator(self):
+    for op in ('==', '!=', '<', '<=', '>', '>='):
+      with self.subTest(op=op):
+        yaml_mapping.validate_generic_expression({'expression': f'age {op} 21'},
+                                                 self.FIELDS,
+                                                 True,
+                                                 'keep')
+
+  def test_accepts_comparisons_without_spaces(self):
+    yaml_mapping.validate_generic_expression({'expression': 'age>=21'},
+                                             self.FIELDS,
+                                             True,
+                                             'keep')
+
+  def test_rejects_non_atomic_sides(self):
+    for expr in ('age + 1 >= 21', 'unknown >= 21', 'age >= 21 >= 1'):
+      with self.subTest(expr=expr):
+        with self.assertRaises(ValueError):
+          yaml_mapping.validate_generic_expression({'expression': expr},
+                                                   self.FIELDS,
+                                                   True,
+                                                   'keep')
+
+  def test_rejects_comparisons_when_not_allowed(self):
+    with self.assertRaises(ValueError):
+      yaml_mapping.validate_generic_expression({'expression': 'age >= 21'},
+                                               self.FIELDS,
+                                               False,
+                                               'field')
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)