diff --git a/CHANGES.md b/CHANGES.md index 59c7ac7b24ba..7ea2cccde299 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -70,6 +70,7 @@ * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). * (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)). +* (Python) Added `equal_to_approx`, an `assert_that` matcher that compares numeric pipeline outputs with a configurable tolerance ([#18028](https://github.com/apache/beam/issues/18028)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/sdks/python/apache_beam/testing/util.py b/sdks/python/apache_beam/testing/util.py index cfe8221b4aa2..bf969d5aa00f 100644 --- a/sdks/python/apache_beam/testing/util.py +++ b/sdks/python/apache_beam/testing/util.py @@ -22,6 +22,8 @@ import collections import glob import io +import math +import numbers import tempfile from typing import Any from typing import Iterable @@ -44,6 +46,7 @@ __all__ = [ 'assert_that', 'equal_to', + 'equal_to_approx', 'equal_to_per_window', 'has_at_least_one', 'is_empty', @@ -231,6 +234,46 @@ def row_namedtuple_equals_fn(expected, actual, fallback_equals_fn=None): return True +def equal_to_approx(expected, rel_tol=1e-09, abs_tol=0.0): + """Matcher used by assert_that that compares numeric elements approximately. + + Behaves similarly to `equal_to` for sequence ordering and membership, but any + real number elements (integers and floats) are compared using `math.isclose` + instead of exact equality. + + Approximate comparisons are also applied to real numbers nested within lists + and tuples. Note that `equal_to`'s advanced handling for Beam `Row` and + `NamedTuple` is not supported here. All other elements are compared using + standard `==` equality. + + Args: + expected: The expected output or sequence to compare against. + rel_tol: The relative tolerance used by `math.isclose` (default: 1e-09). + abs_tol: The absolute tolerance used by `math.isclose` (default: 0.0). + + Example: + assert_that( + pipeline | beam.Create([1.000000001, 2.0]), + equal_to_approx([1.0, 2.0])) + """ + def _approx_equals(expected_element, actual_element): + return _elements_approx_equal( + expected_element, actual_element, rel_tol, abs_tol) + + return equal_to(expected, equals_fn=_approx_equals) + + +def _elements_approx_equal(expected, actual, rel_tol, abs_tol): + if all(isinstance(x, numbers.Real) for x in (expected, actual)): + return math.isclose(expected, actual, rel_tol=rel_tol, abs_tol=abs_tol) + if (isinstance(expected, (list, tuple)) and type(actual) is type(expected) and + len(expected) == len(actual)): + return all( + _elements_approx_equal(e, a, rel_tol, abs_tol) + for e, a in zip(expected, actual)) + return expected == actual + + def matches_all(expected): """Matcher used by assert_that to check a set of matchers. diff --git a/sdks/python/apache_beam/testing/util_test.py b/sdks/python/apache_beam/testing/util_test.py index 12314f4653aa..87f6162ea311 100644 --- a/sdks/python/apache_beam/testing/util_test.py +++ b/sdks/python/apache_beam/testing/util_test.py @@ -29,6 +29,7 @@ from apache_beam.testing.util import TestWindowedValue from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to +from apache_beam.testing.util import equal_to_approx from apache_beam.testing.util import equal_to_per_window from apache_beam.testing.util import is_empty from apache_beam.testing.util import is_not_empty @@ -93,6 +94,49 @@ def test_assert_with_custom_comparator(self): p | Create([1, 2, 3]), equal_to(['1', '2', '3'], equals_fn=lambda e, a: int(e) == int(a))) + def test_equal_to_approx(self): + with TestPipeline() as p: + assert_that( + p | Create([1.0, 2, 3.0]), equal_to_approx([3.0000000001, 2.0, 1.0])) + + def test_equal_to_approx_nested(self): + with TestPipeline() as p: + assert_that( + p | Create([('a', 1.0), ('b', 2.0)]), + equal_to_approx([('b', 2.0000000001), ('a', 1)])) + + def test_equal_to_approx_with_abs_tol(self): + with TestPipeline() as p: + assert_that(p | Create([0.0]), equal_to_approx([1e-10], abs_tol=1e-9)) + + def test_equal_to_approx_fails_outside_tolerance(self): + with self.assertRaises(Exception): + with TestPipeline() as p: + assert_that(p | Create([1.0]), equal_to_approx([1.1])) + + def test_equal_to_approx_nested_list(self): + with TestPipeline() as p: + assert_that( + p | Create([[1.0, 2.0]]), equal_to_approx([[1.0000000001, 2.0]])) + + def test_equal_to_approx_non_numeric(self): + with TestPipeline() as p: + assert_that(p | Create(['a', 'b']), equal_to_approx(['b', 'a'])) + + def test_equal_to_approx_empty(self): + with TestPipeline() as p: + assert_that(p | Create([]), equal_to_approx([])) + + def test_equal_to_approx_with_rel_tol(self): + with TestPipeline() as p: + assert_that( + p | Create([100.0]), equal_to_approx([100.00001], rel_tol=1e-6)) + + def test_equal_to_approx_nested_fails_outside_tolerance(self): + with self.assertRaises(Exception): + with TestPipeline() as p: + assert_that(p | Create([('a', 1.0)]), equal_to_approx([('a', 1.2)])) + def test_reified_value_passes(self): expected = [ TestWindowedValue(v, MIN_TIMESTAMP, [GlobalWindow()])