566 lines
20 KiB
Python
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__]))
|