chore: import upstream snapshot with attribution

python/ray/data/tests/unit/test_distribution_tracker.py

@@ -0,0 +1,167 @@
+import numpy as np
+import pytest
+
+try:
+    import datasketches
+except ImportError:
+    datasketches = None
+
+from ray.data._internal.execution.interfaces.distribution_tracker import (
+    DistributionTracker,
+)
+
+
+def test_empty_tracker_has_zero_moments_and_no_extremes():
+    tracker = DistributionTracker()
+
+    assert tracker.num_samples == 0
+    assert tracker.mean == 0.0
+    assert tracker.variance == 0.0
+    assert tracker.min is None
+    assert tracker.max is None
+
+
+def test_moments_match_numpy_after_adding_samples():
+    tracker = DistributionTracker()
+    samples = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
+    for s in samples:
+        tracker.add_sample(s)
+
+    assert tracker.num_samples == len(samples)
+    assert pytest.approx(tracker.mean) == np.mean(samples)
+    assert pytest.approx(tracker.variance) == np.var(samples, ddof=1)
+    assert pytest.approx(tracker.stddev) == np.std(samples, ddof=1)
+
+
+def test_extremes_track_min_and_max():
+    tracker = DistributionTracker()
+    samples = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
+    for s in samples:
+        tracker.add_sample(s)
+
+    assert tracker.min == 2.0
+    assert tracker.max == 9.0
+
+
+def test_as_dict_contains_all_fields():
+    tracker = DistributionTracker()
+    tracker.add_sample(1.0)
+
+    d = tracker.as_dict()
+    assert set(d.keys()) == {
+        "num_samples",
+        "mean",
+        "variance",
+        "min",
+        "max",
+        "p25",
+        "p50",
+        "p75",
+        "p90",
+        "p95",
+        "p99",
+    }
+
+
+@pytest.mark.skipif(datasketches is None, reason="datasketches not installed")
+def test_percentiles_approximate_expected_quantiles():
+    tracker = DistributionTracker()
+    for i in range(1, 101):
+        tracker.add_sample(float(i))
+
+    assert tracker.p50 is not None and 45 <= tracker.p50 <= 55
+    assert tracker.p90 is not None and 85 <= tracker.p90 <= 95
+    assert tracker.p99 is not None and 95 <= tracker.p99 <= 100
+
+
+def _build(samples):
+    tracker = DistributionTracker()
+    for s in samples:
+        tracker.add_sample(s)
+    return tracker
+
+
+def test_merge_moments_match_numpy_on_concatenation():
+    a = _build([2.0, 4.0, 4.0, 4.0])
+    b = _build([5.0, 5.0, 7.0, 9.0])
+
+    a.merge(b)
+    combined = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
+
+    assert a.num_samples == len(combined)
+    assert pytest.approx(a.mean) == np.mean(combined)
+    assert pytest.approx(a.variance) == np.var(combined, ddof=1)
+    assert a.min == min(combined)
+    assert a.max == max(combined)
+
+
+def test_merge_is_commutative():
+    samples_a = [2.0, 4.0, 4.0, 4.0]
+    samples_b = [5.0, 5.0, 7.0, 9.0]
+
+    ab = _build(samples_a)
+    ab.merge(_build(samples_b))
+    ba = _build(samples_b)
+    ba.merge(_build(samples_a))
+
+    assert ab.num_samples == ba.num_samples
+    assert pytest.approx(ab.mean) == ba.mean
+    assert pytest.approx(ab.variance) == ba.variance
+    assert ab.min == ba.min
+    assert ab.max == ba.max
+
+
+def test_merge_with_empty_other_is_noop():
+    tracker = _build([2.0, 4.0, 6.0])
+
+    tracker.merge(DistributionTracker())
+
+    assert tracker.num_samples == 3
+    assert pytest.approx(tracker.mean) == np.mean([2.0, 4.0, 6.0])
+    assert tracker.min == 2.0
+    assert tracker.max == 6.0
+
+
+def test_merge_self_is_noop():
+    tracker = _build([2.0, 4.0, 6.0])
+
+    tracker.merge(tracker)
+
+    assert tracker.num_samples == 3
+    assert pytest.approx(tracker.mean) == 4.0
+
+
+@pytest.mark.skipif(datasketches is None, reason="datasketches not installed")
+def test_cloudpickle_roundtrip_preserves_sketch():
+    # ``kll_doubles_sketch`` is C++-backed and not natively picklable —
+    # without DistributionTracker's serialize/deserialize hooks, any
+    # Ray Data path that cloudpickles a Dataset (it carries Timers,
+    # which carry DistributionTrackers) fails with
+    # ``TypeError: cannot pickle 'kll_doubles_sketch' object``.
+    import pickle
+
+    import cloudpickle
+
+    tracker = DistributionTracker()
+    for i in range(1, 101):
+        tracker.add_sample(float(i))
+
+    for dumps, loads in [
+        (pickle.dumps, pickle.loads),
+        (cloudpickle.dumps, cloudpickle.loads),
+    ]:
+        restored = loads(dumps(tracker))
+        # Welford moments are exact across the round-trip.
+        assert restored.num_samples == tracker.num_samples
+        assert restored.mean == tracker.mean
+        assert restored.min == tracker.min
+        assert restored.max == tracker.max
+        # The deserialized sketch must still answer quantile queries.
+        assert restored.p50 is not None
+        assert restored.p50 == tracker.p50
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))