Skip to content

[Python] Deflake anomaly transform tests #39669

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
bvolpato wants to merge 2 commits into
base: master
Choose a base branch
from
+70 −66
Draft

[Python] Deflake anomaly transform tests #39669

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b366fc2
[Python] Deflake anomaly transform tests
bvolpato Aug 7, 2026
41bb0af
[Python] Clarify anomaly transform test fixtures
bvolpato Aug 13, 2026
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
136 changes: 70 additions & 66 deletions sdks/python/apache_beam/ml/anomaly/transforms_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,11 @@

import apache_beam as beam
from apache_beam.ml.anomaly.aggregations import AnyVote
from apache_beam.ml.anomaly.base import AnomalyDetector
from apache_beam.ml.anomaly.base import AnomalyPrediction
from apache_beam.ml.anomaly.base import AnomalyResult
from apache_beam.ml.anomaly.base import EnsembleAnomalyDetector
from apache_beam.ml.anomaly.detectors.offline import OfflineDetector
from apache_beam.ml.anomaly.detectors.zscore import ZScore
from apache_beam.ml.anomaly.specifiable import Spec
from apache_beam.ml.anomaly.specifiable import Specifiable
from apache_beam.ml.anomaly.specifiable import _spec_type_to_subspace
Expand Down Expand Up @@ -116,65 +116,58 @@ def _keyed_result_is_equal_to(
return a[0] == b[0] and _unkeyed_result_is_equal_to(a[1], b[1])


@specifiable
class _ValueDetector(AnomalyDetector):
"""Provides order-independent scores for transform orchestration tests.

A score of -1 represents detector warmup and is returned as NaN.
"""
def learn_one(self, unused_x: beam.Row) -> None:
pass

def score_one(self, x: beam.Row) -> float:
value = float(next(iter(x)))
return float('NaN') if value == -1 else value


class TestAnomalyDetection(unittest.TestCase):
class TestData:
unkeyed_input = [
beam.Row(x1=1, x2=4),
beam.Row(x1=2, x2=4),
beam.Row(x1=3, x2=5),
beam.Row(x1=10, x2=4), # outlier in key=1, with respect to x1
beam.Row(x1=2, x2=10), # outlier in key=1, with respect to x2
beam.Row(x1=3, x2=4),
beam.Row(x1=1, x2=4, score_x1=-1.0, score_x2=-1.0),
beam.Row(x1=2, x2=4, score_x1=-1.0, score_x2=-1.0),
beam.Row(x1=3, x2=5, score_x1=2.0, score_x2=0.0),
beam.Row(x1=10, x2=4, score_x1=4.0, score_x2=1.0),
beam.Row(x1=2, x2=10, score_x1=1.0, score_x2=3.0),
beam.Row(x1=3, x2=4, score_x1=0.0, score_x2=1.0),
]
keyed_input = list(zip(itertools.repeat(1), unkeyed_input)) + [
(2, beam.Row(x1=100, x2=5, score_x1=-1.0, score_x2=-1.0))
]
keyed_input = list(zip(itertools.repeat(1),
unkeyed_input)) + [(2, beam.Row(x1=100, x2=5))]

zscore_x1_expected_predictions = [
AnomalyPrediction(
model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3),
AnomalyPrediction(
model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3),
score_x1_expected_predictions = [
AnomalyPrediction(
model_id='zscore_x1',
score=2.1213203435596424,
label=0,
threshold=3),
model_id='value_x1', score=float('NaN'), label=-2, threshold=3),
AnomalyPrediction(
model_id='zscore_x1', score=8.0, label=1, threshold=3),
model_id='value_x1', score=float('NaN'), label=-2, threshold=3),
AnomalyPrediction(model_id='value_x1', score=2.0, label=0, threshold=3),
AnomalyPrediction(model_id='value_x1', score=4.0, label=1, threshold=3),
AnomalyPrediction(model_id='value_x1', score=1.0, label=0, threshold=3),
AnomalyPrediction(model_id='value_x1', score=0.0, label=0, threshold=3),
AnomalyPrediction(
model_id='zscore_x1',
score=0.4898979485566356,
label=0,
threshold=3),
AnomalyPrediction(
model_id='zscore_x1',
score=0.16452254913212455,
label=0,
threshold=3),
AnomalyPrediction(
model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3),
model_id='value_x1', score=float('NaN'), label=-2, threshold=3),
]

zscore_x2_expected_predictions = [
AnomalyPrediction(
model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2),
AnomalyPrediction(
model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2),
AnomalyPrediction(model_id='zscore_x2', score=0, label=0, threshold=2),
score_x2_expected_predictions = [
AnomalyPrediction(
model_id='zscore_x2',
score=0.5773502691896252,
label=0,
threshold=2),
model_id='value_x2', score=float('NaN'), label=-2, threshold=2),
AnomalyPrediction(
model_id='zscore_x2', score=11.5, label=1, threshold=2),
model_id='value_x2', score=float('NaN'), label=-2, threshold=2),
AnomalyPrediction(model_id='value_x2', score=0, label=0, threshold=2),
AnomalyPrediction(model_id='value_x2', score=1.0, label=0, threshold=2),
AnomalyPrediction(model_id='value_x2', score=3.0, label=1, threshold=2),
AnomalyPrediction(model_id='value_x2', score=1.0, label=0, threshold=2),
AnomalyPrediction(
model_id='zscore_x2',
score=0.5368754921931594,
label=0,
threshold=2),
AnomalyPrediction(
model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2),
model_id='value_x2', score=float('NaN'), label=-2, threshold=2),
]

aggregated_expected_predictions = [
Expand All @@ -187,29 +180,29 @@ class TestData:
AnomalyPrediction(model_id="custom", label=-2),
]

keyed_zscore_x1_expected = [
keyed_single_expected = [
(input[0], AnomalyResult(example=input[1], predictions=[decision]))
for input, decision in zip(keyed_input, zscore_x1_expected_predictions)
for input, decision in zip(keyed_input, score_x1_expected_predictions)
]

unkeyed_zscore_x1_expected = [
unkeyed_single_expected = [
AnomalyResult(example=input, predictions=[decision])
for input, decision in zip(
unkeyed_input, zscore_x1_expected_predictions)
unkeyed_input, score_x1_expected_predictions)
]

keyed_ensemble_expected = [(
input[0],
AnomalyResult(example=input[1], predictions=[decision1, decision2]))
for input, decision1, decision2 in zip(
keyed_input, zscore_x1_expected_predictions,
zscore_x2_expected_predictions)]
keyed_input, score_x1_expected_predictions,
score_x2_expected_predictions)]

unkeyed_ensemble_expected = [
AnomalyResult(example=input, predictions=[decision1, decision2])
for input, decision1, decision2 in zip(
unkeyed_input, zscore_x1_expected_predictions,
zscore_x2_expected_predictions)
unkeyed_input, score_x1_expected_predictions,
score_x2_expected_predictions)
]

keyed_ensemble_agg_expected = [
Expand All @@ -225,11 +218,14 @@ class TestData:
]

@parameterized.expand([
(TestData.keyed_input, TestData.keyed_zscore_x1_expected),
(TestData.unkeyed_input, TestData.unkeyed_zscore_x1_expected),
(TestData.keyed_input, TestData.keyed_single_expected),
(TestData.unkeyed_input, TestData.unkeyed_single_expected),
])
def test_one_detector(self, input, expected):
detector = ZScore(features=["x1"], model_id="zscore_x1")
detector = _ValueDetector(
features=["score_x1"],
threshold_criterion=FixedThreshold(3),
model_id="value_x1")
with TestPipeline() as p:
result = (p | beam.Create(input) | AnomalyDetection(detector))

Expand All @@ -244,12 +240,16 @@ def test_one_detector(self, input, expected):
])
def test_multiple_detectors_without_aggregation(self, input, expected):
sub_detectors = []
sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1"))
sub_detectors.append(
ZScore(
features=["x2"],
_ValueDetector(
features=["score_x1"],
threshold_criterion=FixedThreshold(3),
model_id="value_x1"))
sub_detectors.append(
_ValueDetector(
features=["score_x2"],
threshold_criterion=FixedThreshold(2),
model_id="zscore_x2"))
model_id="value_x2"))

with beam.Pipeline() as p:
result = (
Expand All @@ -267,12 +267,16 @@ def test_multiple_detectors_without_aggregation(self, input, expected):
])
def test_multiple_sub_detectors_with_aggregation(self, input, expected):
sub_detectors = []
sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1"))
sub_detectors.append(
ZScore(
features=["x2"],
_ValueDetector(
features=["score_x1"],
threshold_criterion=FixedThreshold(3),
model_id="value_x1"))
sub_detectors.append(
_ValueDetector(
features=["score_x2"],
threshold_criterion=FixedThreshold(2),
model_id="zscore_x2"))
model_id="value_x2"))

with beam.Pipeline() as p:
result = (
Expand Down
Loading