wehub-resource-sync
168 lines
4.5 KiB
Python
168 lines
4.5 KiB
Python
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__]))
|