Files
ray-project--ray/python/ray/data/tests/test_aggregations.py
T
2026-07-13 13:17:40 +08:00

566 lines
20 KiB
Python

from collections import Counter
import numpy as np
import pytest
import ray
from ray.data.aggregate import (
ApproximateQuantile,
ApproximateTopK,
MissingValuePercentage,
Unique,
ZeroPercentage,
)
from ray.data.tests.conftest import * # noqa
from ray.tests.conftest import * # noqa
class TestMissingValuePercentage:
"""Test cases for MissingValuePercentage aggregation."""
def test_missing_value_percentage_basic(self, ray_start_regular_shared_2_cpus):
"""Test basic missing value percentage calculation."""
# Create test data with some null values
data = [
{"id": 1, "value": 10},
{"id": 2, "value": None},
{"id": 3, "value": 30},
{"id": 4, "value": None},
{"id": 5, "value": 50},
]
ds = ray.data.from_items(data)
result = ds.aggregate(MissingValuePercentage(on="value"))
expected = 40.0 # 2 nulls out of 5 total = 40%
assert result["missing_pct(value)"] == expected
def test_missing_value_percentage_no_nulls(self, ray_start_regular_shared_2_cpus):
"""Test missing value percentage with no null values."""
data = [
{"id": 1, "value": 10},
{"id": 2, "value": 20},
{"id": 3, "value": 30},
]
ds = ray.data.from_items(data)
result = ds.aggregate(MissingValuePercentage(on="value"))
expected = 0.0 # 0 nulls out of 3 total = 0%
assert result["missing_pct(value)"] == expected
def test_missing_value_percentage_all_nulls(self, ray_start_regular_shared_2_cpus):
"""Test missing value percentage with all null values."""
data = [
{"id": 1, "value": None},
{"id": 2, "value": None},
{"id": 3, "value": None},
]
ds = ray.data.from_items(data)
result = ds.aggregate(MissingValuePercentage(on="value"))
expected = 100.0 # 3 nulls out of 3 total = 100%
assert result["missing_pct(value)"] == expected
def test_missing_value_percentage_with_nan(self, ray_start_regular_shared_2_cpus):
"""Test missing value percentage with NaN values."""
data = [
{"id": 1, "value": 10.0},
{"id": 2, "value": np.nan},
{"id": 3, "value": None},
{"id": 4, "value": 40.0},
]
ds = ray.data.from_items(data)
result = ds.aggregate(MissingValuePercentage(on="value"))
expected = 50.0 # 2 nulls (NaN + None) out of 4 total = 50%
assert result["missing_pct(value)"] == expected
def test_missing_value_percentage_with_string(
self, ray_start_regular_shared_2_cpus
):
"""Test missing value percentage with string values."""
data = [
{"id": 1, "value": "a"},
{"id": 2, "value": None},
{"id": 3, "value": None},
{"id": 4, "value": "b"},
]
ds = ray.data.from_items(data)
result = ds.aggregate(MissingValuePercentage(on="value"))
expected = 50.0 # 2 None out of 4 total = 50%
assert result["missing_pct(value)"] == expected
def test_missing_value_percentage_custom_alias(
self, ray_start_regular_shared_2_cpus
):
"""Test missing value percentage with custom alias name."""
data = [
{"id": 1, "value": 10},
{"id": 2, "value": None},
]
ds = ray.data.from_items(data)
result = ds.aggregate(MissingValuePercentage(on="value", alias_name="null_pct"))
expected = 50.0 # 1 null out of 2 total = 50%
assert result["null_pct"] == expected
def test_missing_value_percentage_large_dataset(
self, ray_start_regular_shared_2_cpus
):
"""Test missing value percentage with larger dataset."""
# Create a larger dataset with known null percentage
data = []
for i in range(1000):
value = None if i % 10 == 0 else i # 10% null values
data.append({"id": i, "value": value})
ds = ray.data.from_items(data)
result = ds.aggregate(MissingValuePercentage(on="value"))
expected = 10.0 # 100 nulls out of 1000 total = 10%
assert abs(result["missing_pct(value)"] - expected) < 0.01
class TestZeroPercentage:
"""Test cases for ZeroPercentage aggregation."""
def test_zero_percentage_basic(self, ray_start_regular_shared_2_cpus):
"""Test basic zero percentage calculation."""
data = [
{"id": 1, "value": 10},
{"id": 2, "value": 0},
{"id": 3, "value": 30},
{"id": 4, "value": 0},
{"id": 5, "value": 50},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value"))
expected = 40.0 # 2 zeros out of 5 total = 40%
assert result["zero_pct(value)"] == expected
def test_zero_percentage_no_zeros(self, ray_start_regular_shared_2_cpus):
"""Test zero percentage with no zero values."""
data = [
{"id": 1, "value": 10},
{"id": 2, "value": 20},
{"id": 3, "value": 30},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value"))
expected = 0.0 # 0 zeros out of 3 total = 0%
assert result["zero_pct(value)"] == expected
def test_zero_percentage_all_zeros(self, ray_start_regular_shared_2_cpus):
"""Test zero percentage with all zero values."""
data = [
{"id": 1, "value": 0},
{"id": 2, "value": 0},
{"id": 3, "value": 0},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value"))
expected = 100.0 # 3 zeros out of 3 total = 100%
assert result["zero_pct(value)"] == expected
def test_zero_percentage_with_nulls_ignore_nulls_true(
self, ray_start_regular_shared_2_cpus
):
"""Test zero percentage with null values when ignore_nulls=True."""
data = [
{"id": 1, "value": 10},
{"id": 2, "value": 0},
{"id": 3, "value": None},
{"id": 4, "value": 0},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value", ignore_nulls=True))
expected = 66.67 # 2 zeros out of 3 non-null values ≈ 66.67%
assert abs(result["zero_pct(value)"] - expected) < 0.01
def test_zero_percentage_with_nulls_ignore_nulls_false(
self, ray_start_regular_shared_2_cpus
):
"""Test zero percentage with null values when ignore_nulls=False."""
data = [
{"id": 1, "value": 10},
{"id": 2, "value": 0},
{"id": 3, "value": None},
{"id": 4, "value": 0},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value", ignore_nulls=False))
expected = 50.0 # 2 zeros out of 4 total values = 50%
assert result["zero_pct(value)"] == expected
def test_zero_percentage_all_nulls(self, ray_start_regular_shared_2_cpus):
"""Test zero percentage with all null values."""
data = [
{"id": 1, "value": None},
{"id": 2, "value": None},
{"id": 3, "value": None},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value", ignore_nulls=True))
expected = None # No non-null values to calculate percentage
assert result["zero_pct(value)"] == expected
def test_zero_percentage_custom_alias(self, ray_start_regular_shared_2_cpus):
"""Test zero percentage with custom alias name."""
data = [
{"id": 1, "value": 10},
{"id": 2, "value": 0},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value", alias_name="zero_ratio"))
expected = 50.0 # 1 zero out of 2 total = 50%
assert result["zero_ratio"] == expected
def test_zero_percentage_large_dataset(self, ray_start_regular_shared_2_cpus):
"""Test zero percentage with larger dataset."""
# Create a larger dataset with known zero percentage
data = []
for i in range(1000):
value = 0 if i % 5 == 0 else i # 20% zero values
data.append({"id": i, "value": value})
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value"))
expected = 20.0 # 200 zeros out of 1000 total = 20%
assert abs(result["zero_pct(value)"] - expected) < 0.01
def test_zero_percentage_float_zeros(self, ray_start_regular_shared_2_cpus):
"""Test zero percentage with float zero values."""
data = [
{"id": 1, "value": 10.5},
{"id": 2, "value": 0.0},
{"id": 3, "value": 30.7},
{"id": 4, "value": 0.0},
{"id": 5, "value": 50.2},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value"))
expected = 40.0 # 2 zeros out of 5 total = 40%
assert result["zero_pct(value)"] == expected
def test_zero_percentage_negative_values(self, ray_start_regular_shared_2_cpus):
"""Test zero percentage with negative values (zeros should still be counted)."""
data = [
{"id": 1, "value": -10},
{"id": 2, "value": 0},
{"id": 3, "value": 30},
{"id": 4, "value": -5},
{"id": 5, "value": 0},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ZeroPercentage(on="value"))
expected = 40.0 # 2 zeros out of 5 total = 40%
assert result["zero_pct(value)"] == expected
class TestApproximateQuantile:
"""Test cases for ApproximateQuantile aggregation."""
def test_approximate_quantile_basic(self, ray_start_regular_shared_2_cpus):
"""Test basic approximate quantile calculation."""
data = [
{
"id": 1,
"value": 10,
},
{"id": 2, "value": 0},
{"id": 3, "value": 30},
{"id": 4, "value": 0},
{"id": 5, "value": 50},
]
ds = ray.data.from_items(data)
result = ds.aggregate(
ApproximateQuantile(on="value", quantiles=[0.1, 0.5, 0.9])
)
expected = [0.0, 10.0, 50.0]
assert result["approx_quantile(value)"] == expected
def test_approximate_quantile_ignores_nulls(self, ray_start_regular_shared_2_cpus):
data = [
{"id": 1, "value": 5.0},
{"id": 2, "value": None},
{"id": 3, "value": 15.0},
{"id": 4, "value": None},
{"id": 5, "value": 25.0},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ApproximateQuantile(on="value", quantiles=[0.5]))
assert result["approx_quantile(value)"] == [15.0]
def test_approximate_quantile_custom_alias(self, ray_start_regular_shared_2_cpus):
data = [
{"id": 1, "value": 1.0},
{"id": 2, "value": 3.0},
{"id": 3, "value": 5.0},
{"id": 4, "value": 7.0},
{"id": 5, "value": 9.0},
]
ds = ray.data.from_items(data)
quantiles = [0.0, 1.0]
result = ds.aggregate(
ApproximateQuantile(
on="value", quantiles=quantiles, alias_name="value_range"
)
)
assert result["value_range"] == [1.0, 9.0]
assert len(result["value_range"]) == len(quantiles)
def test_approximate_quantile_groupby(self, ray_start_regular_shared_2_cpus):
data = [
{"group": "A", "value": 1.0},
{"group": "A", "value": 2.0},
{"group": "A", "value": 3.0},
{"group": "B", "value": 10.0},
{"group": "B", "value": 20.0},
{"group": "B", "value": 30.0},
]
ds = ray.data.from_items(data)
result = (
ds.groupby("group")
.aggregate(ApproximateQuantile(on="value", quantiles=[0.5]))
.take_all()
)
result_by_group = {
row["group"]: row["approx_quantile(value)"] for row in result
}
assert result_by_group["A"] == [2.0]
assert result_by_group["B"] == [20.0]
class TestApproximateTopK:
"""Test cases for ApproximateTopK aggregation."""
def test_approximate_topk_ignores_nulls(self, ray_start_regular_shared_2_cpus):
"""Test that null values are ignored."""
data = [
*[{"word": "apple"} for _ in range(5)],
*[{"word": None} for _ in range(10)],
*[{"word": "banana"} for _ in range(3)],
*[{"word": "cherry"} for _ in range(2)],
]
ds = ray.data.from_items(data)
result = ds.aggregate(ApproximateTopK(on="word", k=2))
assert result["approx_topk(word)"] == [
{"word": "apple", "count": 5},
{"word": "banana", "count": 3},
]
def test_approximate_topk_custom_alias(self, ray_start_regular_shared_2_cpus):
"""Test approximate top k with custom alias."""
data = [
*[{"item": "x"} for _ in range(3)],
*[{"item": "y"} for _ in range(2)],
*[{"item": "z"} for _ in range(1)],
]
ds = ray.data.from_items(data)
result = ds.aggregate(ApproximateTopK(on="item", k=2, alias_name="top_items"))
assert "top_items" in result
assert result["top_items"] == [
{"item": "x", "count": 3},
{"item": "y", "count": 2},
]
def test_approximate_topk_groupby(self, ray_start_regular_shared_2_cpus):
"""Test approximate top k with groupby."""
data = [
*[{"category": "A", "item": "apple"} for _ in range(5)],
*[{"category": "A", "item": "banana"} for _ in range(3)],
*[{"category": "B", "item": "cherry"} for _ in range(4)],
*[{"category": "B", "item": "date"} for _ in range(2)],
]
ds = ray.data.from_items(data)
result = (
ds.groupby("category").aggregate(ApproximateTopK(on="item", k=1)).take_all()
)
result_by_category = {
row["category"]: row["approx_topk(item)"] for row in result
}
assert result_by_category["A"] == [{"item": "apple", "count": 5}]
assert result_by_category["B"] == [{"item": "cherry", "count": 4}]
def test_approximate_topk_all_unique(self, ray_start_regular_shared_2_cpus):
"""Test approximate top k when all items are unique."""
data = [{"id": f"item_{i}"} for i in range(10)]
ds = ray.data.from_items(data)
result = ds.aggregate(ApproximateTopK(on="id", k=3))
# All items have count 1, so we should get exactly 3 items
assert len(result["approx_topk(id)"]) == 3
for item in result["approx_topk(id)"]:
assert item["count"] == 1
def test_approximate_topk_fewer_items_than_k(self, ray_start_regular_shared_2_cpus):
"""Test approximate top k when dataset has fewer unique items than k."""
data = [
{"id": "a"},
{"id": "b"},
]
ds = ray.data.from_items(data)
result = ds.aggregate(ApproximateTopK(on="id", k=5))
# Should only return 2 items since that's all we have
assert len(result["approx_topk(id)"]) == 2
def test_approximate_topk_different_log_capacity(
self, ray_start_regular_shared_2_cpus
):
"""Test that different log_capacity values still produce correct top k."""
data = [
*[{"id": "frequent"} for _ in range(100)],
*[{"id": "common"} for _ in range(50)],
*[{"id": f"rare_{i}"} for i in range(50)], # 50 unique rare items
]
ds = ray.data.from_items(data)
# Test with smaller log_capacity
result_small = ds.aggregate(ApproximateTopK(on="id", k=2, log_capacity=10))
# Test with larger log_capacity
result_large = ds.aggregate(ApproximateTopK(on="id", k=2, log_capacity=15))
# Both should correctly identify the top 2
for result in [result_small, result_large]:
assert result["approx_topk(id)"][0] == {"id": "frequent", "count": 100}
assert result["approx_topk(id)"][1] == {"id": "common", "count": 50}
@pytest.mark.parametrize(
("data", "expected1", "expected2"),
[
(
[{"id": 1}, {"id": 1}, {"id": 2}],
{"id": 1, "count": 2},
{"id": 2, "count": 1},
),
(
[{"id": [1, 2, 3]}, {"id": [1, 2, 3]}, {"id": [1, 2]}],
{"id": [1, 2, 3], "count": 2},
{"id": [1, 2], "count": 1},
),
],
)
def test_approximate_topk_non_string_datatype(
self, data, expected1, expected2, ray_start_regular_shared_2_cpus
):
"""Test that ApproximateTopK works with non-string type elements."""
ds = ray.data.from_items(data)
result = ds.aggregate(ApproximateTopK(on="id", k=2, log_capacity=3))
assert result["approx_topk(id)"][0] == expected1
assert result["approx_topk(id)"][1] == expected2
def test_approximate_topk_encode_lists(self, ray_start_regular_shared_2_cpus):
"""Test ApproximateTopK list encode feature."""
data = [{"id": [1, 1, 1]}, {"id": [2, 2]}, {"id": [3]}]
ds = ray.data.from_items(data)
result = ds.aggregate(
ApproximateTopK(on="id", k=4, log_capacity=10, encode_lists=True)
)
assert result["approx_topk(id)"][0] == {"id": 1, "count": 3}
assert result["approx_topk(id)"][1] == {"id": 2, "count": 2}
assert result["approx_topk(id)"][2] == {"id": 3, "count": 1}
class TestUnique:
"""Test cases for Unique aggregation."""
def test_unique_basic(self, ray_start_regular_shared_2_cpus):
"""Test basic Unique aggregation."""
data = [{"id": "a"}, {"id": "b"}, {"id": "b"}, {"id": None}]
ds = ray.data.from_items(data)
result = ds.aggregate(Unique(on="id", ignore_nulls=False))
assert Counter(result["unique(id)"]) == Counter(["a", "b", None])
def test_unique_ignores_nulls(self, ray_start_regular_shared_2_cpus):
"""Test Unique properly ignores nulls."""
data = [{"id": "a"}, {"id": None}, {"id": "b"}, {"id": "b"}, {"id": None}]
ds = ray.data.from_items(data)
result = ds.aggregate(Unique(on="id", ignore_nulls=True))
assert Counter(result["unique(id)"]) == Counter(["a", "b"])
def test_unique_custom_alias(self, ray_start_regular_shared_2_cpus):
"""Test Unique with custom alias."""
data = [{"id": "a"}, {"id": "b"}, {"id": "b"}]
ds = ray.data.from_items(data)
result = ds.aggregate(Unique(on="id", alias_name="custom"))
assert sorted(result["custom"]) == ["a", "b"]
def test_unique_list_datatype(self, ray_start_regular_shared_2_cpus):
"""Test Unique works with non-hashable types like list."""
data = [
{"id": ["a", "b", "c"]},
{"id": ["a", "b", "c"]},
{"id": ["a", "b", "c"]},
]
ds = ray.data.from_items(data)
result = ds.aggregate(Unique(on="id"))
assert result["unique(id)"][0] == ["a", "b", "c"]
def test_unique_encode_lists(self, ray_start_regular_shared_2_cpus):
"""Test Unique works when encode_lists is True."""
data = [{"id": ["a", "b", "c"]}, {"id": ["a", "a", "a", "b", None]}]
ds = ray.data.from_items(data)
result = ds.aggregate(Unique(on="id", encode_lists=True, ignore_nulls=False))
answer = ["a", "b", "c", None]
assert Counter(result["unique(id)"]) == Counter(answer)
def test_unique_encode_lists_ignores_nulls(self, ray_start_regular_shared_2_cpus):
"""Test Unique will drop null values when encode_lists is True."""
data = [{"id": ["a", "b", "c"]}, {"id": ["a", "a", "a", "b", None]}]
ds = ray.data.from_items(data)
result = ds.aggregate(Unique(on="id", encode_lists=True, ignore_nulls=True))
answer = ["a", "b", "c"]
assert Counter(result["unique(id)"]) == Counter(answer)
if __name__ == "__main__":
import sys
sys.exit(pytest.main(["-v", __file__]))