wehub-resource-sync
1013 lines
36 KiB
Python
1013 lines
36 KiB
Python
"""Unit tests for GPU hash aggregation with ShuffleStrategy.GPU_SHUFFLE.
|
|
|
|
The planning tests do not require GPUs. Tests marked ``gpu`` exercise cuDF and
|
|
RAPIDS MPF paths on actual GPU hardware.
|
|
"""
|
|
|
|
from typing import Any, List, Optional, Tuple, cast
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
import ray.data._internal.gpu_shuffle.hash_aggregate as hash_aggregate
|
|
from ray.actor import ActorClass, ActorHandle
|
|
from ray.data._internal.execution.interfaces import ExecutionResources, PhysicalOperator
|
|
from ray.data._internal.gpu_shuffle.hash_aggregate import (
|
|
GPUAggregateFn,
|
|
GPUAggregationPlan,
|
|
GPUHashAggregateActor,
|
|
GPUHashAggregateOperator,
|
|
build_gpu_aggregation_plan,
|
|
)
|
|
from ray.data._internal.logical.interfaces import LogicalOperator
|
|
from ray.data._internal.logical.operators import Aggregate
|
|
from ray.data._internal.planner.plan_all_to_all_op import plan_all_to_all_op
|
|
from ray.data.aggregate import AggregateFnV2, AsList, Count, Max, Mean, Min, Sum
|
|
from ray.data.block import BlockMetadata
|
|
from ray.data.context import DataContext, ShuffleStrategy
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_input_op_mock(num_blocks=None, size_bytes=None):
|
|
"""Return a minimal PhysicalOperator mock.
|
|
|
|
The mock is compatible with HashShufflingOperatorBase.
|
|
"""
|
|
logical_mock = MagicMock(LogicalOperator)
|
|
logical_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=size_bytes,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
logical_mock.estimated_num_outputs.return_value = num_blocks
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
return op_mock
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU hash aggregate planning
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGPUHashAggregatePlanning:
|
|
def _make_aggregate_op(
|
|
self,
|
|
aggs,
|
|
key="user_id",
|
|
num_partitions=8,
|
|
input_schema=pa.schema([("user_id", pa.int64()), ("value", pa.int64())]),
|
|
):
|
|
input_op = MagicMock(LogicalOperator)
|
|
input_op.infer_schema.return_value = input_schema
|
|
return Aggregate(
|
|
key=key,
|
|
aggs=list(aggs),
|
|
input_dependencies=[input_op],
|
|
num_partitions=num_partitions,
|
|
)
|
|
|
|
def test_builtin_aggregation_plan_supported(self):
|
|
schema = pa.schema([("user_id", pa.int64()), ("value", pa.int64())])
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",),
|
|
(
|
|
Count(),
|
|
Count(on="value", ignore_nulls=True),
|
|
Sum("value"),
|
|
Min("value"),
|
|
Max("value"),
|
|
Mean("value"),
|
|
),
|
|
input_schema=schema,
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
assert plan.shuffle_key_columns == ("user_id",)
|
|
assert plan.output_names == (
|
|
"count()",
|
|
"count(value)",
|
|
"sum(value)",
|
|
"min(value)",
|
|
"max(value)",
|
|
"mean(value)",
|
|
)
|
|
assert tuple(type(agg).__name__ for agg in plan._gpu_aggregates) == (
|
|
"GPUCount",
|
|
"GPUCount",
|
|
"GPUSum",
|
|
"GPUMin",
|
|
"GPUMax",
|
|
"GPUMean",
|
|
)
|
|
for gpu_agg in plan._gpu_aggregates:
|
|
assert not isinstance(gpu_agg, AggregateFnV2)
|
|
assert not hasattr(gpu_agg, "aggregate_block")
|
|
assert not hasattr(gpu_agg, "combine")
|
|
assert not hasattr(gpu_agg, "finalize")
|
|
assert "user_id" in plan.required_columns
|
|
assert "value" in plan.required_columns
|
|
|
|
def test_unsupported_aggregation_plan_rejected(self):
|
|
schema = pa.schema([("user_id", pa.int64())])
|
|
unsupported_result = build_gpu_aggregation_plan(
|
|
("user_id",), (AsList("value"),), input_schema=schema
|
|
)
|
|
assert unsupported_result == "AsList is not supported by GPU aggregation."
|
|
|
|
missing_schema_result = build_gpu_aggregation_plan(
|
|
("user_id",), (Sum("value"),)
|
|
)
|
|
assert missing_schema_result == (
|
|
"missing input schema for key column(s): user_id."
|
|
)
|
|
|
|
def test_custom_gpu_aggregation_plan_supported_without_aggregate_fn_v2(self):
|
|
class _CustomGPUAggregate(GPUAggregateFn):
|
|
def __init__(self) -> None:
|
|
super().__init__(
|
|
"custom(value)",
|
|
on="value",
|
|
ignore_nulls=True,
|
|
accumulators=("acc",),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Any = None,
|
|
) -> Any:
|
|
raise NotImplementedError
|
|
|
|
def final_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
output_name: str,
|
|
) -> Any:
|
|
raise NotImplementedError
|
|
|
|
gpu_agg = _CustomGPUAggregate()
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",),
|
|
(gpu_agg,),
|
|
input_schema=pa.schema([("user_id", pa.int64()), ("value", pa.int64())]),
|
|
)
|
|
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
assert plan._gpu_aggregates == (gpu_agg,)
|
|
assert plan.output_names == ("custom(value)",)
|
|
assert plan.required_columns == ("user_id", "value")
|
|
assert not isinstance(gpu_agg, AggregateFnV2)
|
|
assert not hasattr(gpu_agg, "aggregate_block")
|
|
assert not hasattr(gpu_agg, "combine")
|
|
assert not hasattr(gpu_agg, "finalize")
|
|
for method_name in (
|
|
"_empty_global_partial_values",
|
|
"_partial_accumulator_dtypes",
|
|
"_final_arrow_types",
|
|
"_final_cudf_dtypes",
|
|
):
|
|
assert hasattr(gpu_agg, method_name)
|
|
|
|
def test_custom_gpu_aggregation_plan_supported_without_input_schema(self):
|
|
class _CustomGPUAggregate(GPUAggregateFn):
|
|
def __init__(self) -> None:
|
|
super().__init__(
|
|
"custom(value)",
|
|
on="value",
|
|
ignore_nulls=True,
|
|
accumulators=("acc",),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Any = None,
|
|
) -> Any:
|
|
raise NotImplementedError
|
|
|
|
def final_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
output_name: str,
|
|
) -> Any:
|
|
raise NotImplementedError
|
|
|
|
gpu_agg = _CustomGPUAggregate()
|
|
plan = build_gpu_aggregation_plan(("user_id",), (gpu_agg,), input_schema=None)
|
|
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
assert plan._gpu_aggregates == (gpu_agg,)
|
|
assert plan.output_names == ("custom(value)",)
|
|
assert plan.required_columns == ("user_id", "value")
|
|
|
|
def test_gpu_shuffle_routes_custom_gpu_aggregate_without_input_schema(self):
|
|
class _CustomGPUAggregate(GPUAggregateFn):
|
|
def __init__(self) -> None:
|
|
super().__init__(
|
|
"custom(value)",
|
|
on="value",
|
|
ignore_nulls=True,
|
|
accumulators=("acc",),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Any = None,
|
|
) -> Any:
|
|
raise NotImplementedError
|
|
|
|
def final_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
output_name: str,
|
|
) -> Any:
|
|
raise NotImplementedError
|
|
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
|
|
logical_op = self._make_aggregate_op([_CustomGPUAggregate()], input_schema=None)
|
|
input_physical_op = _make_input_op_mock()
|
|
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert isinstance(op, GPUHashAggregateOperator)
|
|
|
|
def test_gpu_shuffle_routes_supported_aggregate_to_gpu_operator(self):
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
|
|
logical_op = self._make_aggregate_op([Count(), Sum("value")])
|
|
input_physical_op = _make_input_op_mock()
|
|
original_build_plan = hash_aggregate.build_gpu_aggregation_plan
|
|
built_plans: List[GPUAggregationPlan] = []
|
|
|
|
def _build_plan_once(*args: Any, **kwargs: Any) -> GPUAggregationPlan:
|
|
result = original_build_plan(*args, **kwargs)
|
|
assert isinstance(result, GPUAggregationPlan), result
|
|
built_plans.append(result)
|
|
return result
|
|
|
|
with patch(
|
|
"ray.data._internal.gpu_shuffle.hash_aggregate.build_gpu_aggregation_plan",
|
|
side_effect=_build_plan_once,
|
|
) as mock_build_plan:
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert isinstance(op, GPUHashAggregateOperator)
|
|
assert op._num_partitions == 8
|
|
assert "GPUHashAggregate" in op.name
|
|
mock_build_plan.assert_called_once()
|
|
assert op._aggregation_plan is built_plans[0]
|
|
|
|
def test_gpu_hash_aggregate_operator_uses_prebuilt_plan(self):
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
input_physical_op = _make_input_op_mock()
|
|
schema = pa.schema([("user_id", pa.int64()), ("value", pa.int64())])
|
|
aggregation_plan = build_gpu_aggregation_plan(
|
|
("user_id",), (Count(), Sum("value")), input_schema=schema
|
|
)
|
|
assert isinstance(aggregation_plan, GPUAggregationPlan), aggregation_plan
|
|
|
|
with patch(
|
|
"ray.data._internal.gpu_shuffle.hash_shuffle.GPURankPool"
|
|
) as mock_default_pool:
|
|
op = GPUHashAggregateOperator(
|
|
ctx,
|
|
input_physical_op,
|
|
key_columns=("user_id",),
|
|
aggregation_plan=aggregation_plan,
|
|
num_partitions=8,
|
|
)
|
|
|
|
mock_default_pool.assert_not_called()
|
|
assert op._aggregation_plan is aggregation_plan
|
|
assert op._rank_pool.nranks == 4
|
|
|
|
def test_gpu_shuffle_unsupported_aggregate_falls_back_to_cpu_hash_aggregate(self):
|
|
from ray.data._internal.execution.operators.hash_aggregate import (
|
|
HashAggregateOperator,
|
|
)
|
|
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
|
|
logical_op = self._make_aggregate_op([AsList("value")])
|
|
input_physical_op = _make_input_op_mock(num_blocks=8, size_bytes=1024)
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle"
|
|
"._get_total_cluster_resources",
|
|
return_value=ExecutionResources(cpu=4, memory=1024 * 1024 * 1024),
|
|
), patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.put",
|
|
return_value=MagicMock(),
|
|
):
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert isinstance(op, HashAggregateOperator)
|
|
|
|
def test_gpu_shuffle_missing_key_schema_falls_back_to_cpu_hash_aggregate(self):
|
|
from ray.data._internal.execution.operators.hash_aggregate import (
|
|
HashAggregateOperator,
|
|
)
|
|
|
|
ctx = DataContext()
|
|
ctx.gpu_shuffle_num_actors = 4
|
|
ctx._shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE
|
|
|
|
logical_op = self._make_aggregate_op([Sum("value")], input_schema=None)
|
|
input_physical_op = _make_input_op_mock(num_blocks=8, size_bytes=1024)
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle"
|
|
"._get_total_cluster_resources",
|
|
return_value=ExecutionResources(cpu=4, memory=1024 * 1024 * 1024),
|
|
), patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.put",
|
|
return_value=MagicMock(),
|
|
):
|
|
op = plan_all_to_all_op(logical_op, [input_physical_op], ctx)
|
|
|
|
assert isinstance(op, HashAggregateOperator)
|
|
|
|
def test_global_aggregate_uses_synthetic_shuffle_key(self):
|
|
plan = build_gpu_aggregation_plan(tuple(), (Count(), Sum("value")))
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
assert plan.shuffle_key_columns == (hash_aggregate._GLOBAL_AGGREGATE_KEY,)
|
|
|
|
def test_required_columns_excludes_unused_input_columns(self):
|
|
table = pa.table(
|
|
{
|
|
"user_id": pa.array([0, 0, 1], type=pa.int64()),
|
|
"small": pa.array([1, 2, 3], type=pa.int64()),
|
|
"unused": pa.array([[1], [2], [3]], type=pa.list_(pa.int64())),
|
|
}
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",), (Count(), Sum("small")), input_schema=table.schema
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
assert plan.required_columns == ("user_id", "small")
|
|
|
|
projected = table.select(list(plan.required_columns))
|
|
assert projected.column_names == ["user_id", "small"]
|
|
|
|
def test_global_count_required_columns_empty(self):
|
|
table = pa.table({"unused": pa.array([10, 20, 30], type=pa.int64())})
|
|
plan = build_gpu_aggregation_plan(
|
|
tuple(), (Count(),), input_schema=table.schema
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
assert plan.required_columns == ()
|
|
|
|
def test_merge_input_schema_unifies_value_dtype_across_blocks(self):
|
|
block1 = pa.table({"value": pa.array([1], type=pa.int8())})
|
|
block2 = pa.table({"value": pa.array([2], type=pa.int32())})
|
|
plan = build_gpu_aggregation_plan(tuple(), (Sum("value"),))
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
runtime_schema = None
|
|
runtime_schema = plan.merge_input_schema(runtime_schema, block1.schema)
|
|
runtime_schema = plan.merge_input_schema(runtime_schema, block2.schema)
|
|
|
|
assert runtime_schema is not None
|
|
assert runtime_schema.field("value").type == pa.int32()
|
|
|
|
def test_merge_input_schema_uses_logical_schema_for_shuffle_keys(self):
|
|
logical_schema = pa.schema(
|
|
[
|
|
("user_id", pa.int32()),
|
|
("value", pa.int64()),
|
|
]
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",), (Sum("value"),), input_schema=logical_schema
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
runtime_schema = plan.merge_input_schema(
|
|
None, pa.schema([("user_id", pa.int8()), ("value", pa.int64())])
|
|
)
|
|
runtime_schema = plan.merge_input_schema(
|
|
runtime_schema,
|
|
pa.schema([("user_id", pa.int32()), ("value", pa.int64())]),
|
|
)
|
|
|
|
assert runtime_schema is not None
|
|
assert runtime_schema.field("user_id").type == pa.int32()
|
|
assert runtime_schema.field("value").type == pa.int64()
|
|
|
|
def test_normalize_output_arrow_null_schema(self):
|
|
null_schema = pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("value", pa.null()),
|
|
]
|
|
)
|
|
null_plan = build_gpu_aggregation_plan(
|
|
("user_id",), (Sum("value"),), input_schema=null_schema
|
|
)
|
|
assert isinstance(null_plan, GPUAggregationPlan), null_plan
|
|
output_table = pa.table(
|
|
{
|
|
"user_id": pa.array([0, 1, 2], type=pa.int64()),
|
|
"sum(value)": pa.array([None, None, None], type=pa.int64()),
|
|
}
|
|
)
|
|
|
|
assert null_plan.normalize_output_arrow(output_table).schema == pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("sum(value)", pa.null()),
|
|
]
|
|
)
|
|
|
|
def test_normalize_output_arrow_preserves_values_when_runtime_schema_upgrades(
|
|
self,
|
|
):
|
|
null_schema = pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("value", pa.null()),
|
|
]
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",), (Sum("value"),), input_schema=null_schema
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
runtime_schema = plan.merge_input_schema(
|
|
null_schema,
|
|
pa.schema([("user_id", pa.int64()), ("value", pa.int64())]),
|
|
)
|
|
output_table = pa.table(
|
|
{
|
|
"user_id": pa.array([0, 1], type=pa.int64()),
|
|
"sum(value)": pa.array([10, 20], type=pa.int64()),
|
|
}
|
|
)
|
|
|
|
normalized = plan.normalize_output_arrow(
|
|
output_table, input_schema=runtime_schema
|
|
)
|
|
assert normalized.column("sum(value)").to_pylist() == [10, 20]
|
|
assert normalized.schema.field("sum(value)").type == pa.int64()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU fixtures — shared by real-GPU aggregate tests below
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _num_cluster_gpus() -> int:
|
|
"""Return the number of GPUs in the Ray cluster (0 if Ray not initialised)."""
|
|
if not ray.is_initialized():
|
|
return 0
|
|
return int(ray.cluster_resources().get("GPU", 0))
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ray_with_gpu():
|
|
"""Skip the test if GPU packages or GPU hardware are absent."""
|
|
pytest.importorskip("cudf", reason="cudf (GPU DataFrame library) not installed")
|
|
pytest.importorskip("rapidsmpf", reason="rapidsmpf not installed")
|
|
|
|
if not ray.is_initialized():
|
|
ray.init()
|
|
|
|
num_gpus = _num_cluster_gpus()
|
|
if num_gpus < 1:
|
|
pytest.skip("No GPU resources found in the Ray cluster")
|
|
|
|
yield num_gpus
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUAggregationPlan — real cuDF execution (conditional)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.gpu
|
|
class TestGPUAggregationPlanReal:
|
|
"""Exercises aggregation plan methods with real cuDF on GPU hardware."""
|
|
|
|
def test_sum_and_mean_schema_source_dtypes_use_cudf_accumulators(
|
|
self, ray_with_gpu
|
|
):
|
|
import cudf
|
|
|
|
schema = pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("small", pa.int8()),
|
|
("flag", pa.bool_()),
|
|
("ratio", pa.float32()),
|
|
]
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",),
|
|
(Sum("small"), Mean("flag"), Sum("ratio")),
|
|
input_schema=schema,
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
df = cudf.DataFrame(
|
|
{
|
|
"user_id": [0, 0, 1],
|
|
"small": np.array([100, 100, 1], dtype=np.int8),
|
|
"flag": [True, True, False],
|
|
"ratio": np.array([1.5, 2.5, 3.5], dtype=np.float32),
|
|
}
|
|
)
|
|
partial = plan.partial_aggregate(df, input_schema=schema)
|
|
|
|
sum_col = plan.accumulator_columns[0]
|
|
mean_sum_col = plan.accumulator_columns[1]
|
|
float_sum_col = plan.accumulator_columns[4]
|
|
group_zero = partial[partial["user_id"] == 0].iloc[0]
|
|
|
|
assert str(partial[sum_col].dtype) == "int64"
|
|
assert str(partial[mean_sum_col].dtype) == "int64"
|
|
assert str(partial[float_sum_col].dtype) == "float64"
|
|
assert group_zero[sum_col] == 200
|
|
assert group_zero[mean_sum_col] == 2
|
|
assert group_zero[float_sum_col] == pytest.approx(4.0)
|
|
|
|
def test_empty_final_aggregate_preserves_output_dtypes(self, ray_with_gpu):
|
|
import cudf
|
|
|
|
schema = pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("value", pa.int8()),
|
|
]
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",),
|
|
(Count(), Sum("value"), Min("value"), Mean("value")),
|
|
input_schema=schema,
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
result = plan.final_aggregate(cudf.DataFrame())
|
|
result_table = result.to_arrow(preserve_index=False)
|
|
expected_schema = pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("count()", pa.int64()),
|
|
("sum(value)", pa.int64()),
|
|
("min(value)", pa.int64()),
|
|
("mean(value)", pa.float64()),
|
|
]
|
|
)
|
|
assert result_table.schema.equals(expected_schema, check_metadata=False)
|
|
|
|
global_plan = build_gpu_aggregation_plan(
|
|
tuple(), (Count(),), input_schema=schema
|
|
)
|
|
assert isinstance(global_plan, GPUAggregationPlan), global_plan
|
|
|
|
global_result = global_plan.final_aggregate(cudf.DataFrame())
|
|
global_table = global_result.to_arrow(preserve_index=False)
|
|
assert global_table.schema.equals(
|
|
pa.schema([("count()", pa.int64())]), check_metadata=False
|
|
)
|
|
|
|
runtime_plan = build_gpu_aggregation_plan(
|
|
("user_id",),
|
|
(Count(), Sum("value")),
|
|
input_schema=pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("value", pa.int8()),
|
|
]
|
|
),
|
|
)
|
|
assert isinstance(runtime_plan, GPUAggregationPlan), runtime_plan
|
|
|
|
runtime_result = runtime_plan.final_aggregate(
|
|
cudf.DataFrame(),
|
|
input_schema=pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("value", pa.int8()),
|
|
]
|
|
),
|
|
)
|
|
runtime_table = runtime_result.to_arrow(preserve_index=False)
|
|
assert runtime_table.schema.equals(
|
|
pa.schema(
|
|
[
|
|
("user_id", pa.int64()),
|
|
("count()", pa.int64()),
|
|
("sum(value)", pa.int64()),
|
|
]
|
|
),
|
|
check_metadata=False,
|
|
)
|
|
|
|
def test_custom_gpu_aggregate_fn_receives_input_schema(self, ray_with_gpu):
|
|
import cudf
|
|
|
|
class _CustomGPUAggregate(GPUAggregateFn):
|
|
def __init__(self) -> None:
|
|
super().__init__(
|
|
"custom(value)",
|
|
on="value",
|
|
ignore_nulls=True,
|
|
accumulators=("acc",),
|
|
)
|
|
self.seen_schema: Optional[pa.Schema] = None
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Any = None,
|
|
) -> cudf.DataFrame:
|
|
self.seen_schema = input_schema
|
|
acc_col = accumulator_columns[0]
|
|
return df[[key_columns[0], "value"]].rename(columns={"value": acc_col})
|
|
|
|
def final_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
output_name: str,
|
|
) -> cudf.DataFrame:
|
|
acc_col = accumulator_columns[0]
|
|
return df[[key_columns[0], acc_col]].rename(
|
|
columns={acc_col: output_name}
|
|
)
|
|
|
|
schema = pa.schema([("user_id", pa.int64()), ("value", pa.int32())])
|
|
gpu_agg = _CustomGPUAggregate()
|
|
plan = build_gpu_aggregation_plan(("user_id",), (gpu_agg,), input_schema=schema)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
partial = plan.partial_aggregate(
|
|
cudf.DataFrame({"user_id": [1], "value": np.array([2], dtype=np.int32)}),
|
|
input_schema=schema,
|
|
)
|
|
result = plan.final_aggregate(partial, input_schema=schema)
|
|
|
|
assert gpu_agg.seen_schema is schema
|
|
assert partial.to_pandas().to_dict("list") == {
|
|
"user_id": [1],
|
|
plan.accumulator_columns[0]: [2],
|
|
}
|
|
assert str(partial[plan.accumulator_columns[0]].dtype) == "int32"
|
|
assert result.to_pandas().to_dict("list") == {
|
|
"user_id": [1],
|
|
"custom(value)": [2],
|
|
}
|
|
assert str(result["custom(value)"].dtype) == "int32"
|
|
|
|
def test_null_reductions_preserve_groups_and_accumulator_dtypes(self, ray_with_gpu):
|
|
import cudf
|
|
|
|
schema = pa.schema([("user_id", pa.int64())])
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",), (Sum("value"),), input_schema=schema
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
acc_col = plan.accumulator_columns[0]
|
|
|
|
df = cudf.DataFrame(
|
|
{
|
|
"user_id": [0, 1, 2, 0],
|
|
"value": cudf.Series([None, None, None, None], dtype="int64"),
|
|
}
|
|
)
|
|
partial = plan.partial_aggregate(df)
|
|
assert str(partial[acc_col].dtype) == "int64"
|
|
|
|
result = (
|
|
plan.final_aggregate(partial)
|
|
.to_pandas()
|
|
.sort_values("user_id")
|
|
.reset_index(drop=True)
|
|
)
|
|
assert result["user_id"].tolist() == [0, 1, 2]
|
|
assert result["sum(value)"].isna().all()
|
|
|
|
count_plan = build_gpu_aggregation_plan(
|
|
("user_id",),
|
|
(Count("value", ignore_nulls=True),),
|
|
input_schema=schema,
|
|
)
|
|
assert isinstance(count_plan, GPUAggregationPlan), count_plan
|
|
|
|
count_partial = count_plan.partial_aggregate(df)
|
|
count_result = (
|
|
count_plan.final_aggregate(count_partial)
|
|
.to_pandas()
|
|
.sort_values("user_id")
|
|
.reset_index(drop=True)
|
|
)
|
|
assert count_result.to_dict("records") == [
|
|
{"user_id": 0, "count(value)": 0},
|
|
{"user_id": 1, "count(value)": 0},
|
|
{"user_id": 2, "count(value)": 0},
|
|
]
|
|
|
|
def test_partial_aggregate_widens_shuffle_key_to_logical_schema(self, ray_with_gpu):
|
|
import cudf
|
|
|
|
logical_schema = pa.schema(
|
|
[
|
|
("user_id", pa.int32()),
|
|
("value", pa.int64()),
|
|
]
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("user_id",), (Sum("value"),), input_schema=logical_schema
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
df = cudf.DataFrame(
|
|
{
|
|
"user_id": np.array([1], dtype=np.int8),
|
|
"value": np.array([1], dtype=np.int64),
|
|
}
|
|
)
|
|
partial = plan.partial_aggregate(
|
|
df,
|
|
input_schema=pa.schema([("user_id", pa.int8()), ("value", pa.int64())]),
|
|
)
|
|
|
|
assert partial["user_id"].iloc[0] == 1
|
|
assert str(partial["user_id"].dtype) == "int32"
|
|
|
|
def test_global_count_partial_aggregate(self, ray_with_gpu):
|
|
import cudf
|
|
|
|
plan = build_gpu_aggregation_plan(tuple(), (Count(),))
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
df = cudf.DataFrame(index=range(3))
|
|
partial = plan.partial_aggregate(df)
|
|
assert partial[plan.accumulator_columns[0]].iloc[0] == 3
|
|
|
|
def test_partial_aggregate_normalizes_null_key_column(self, ray_with_gpu):
|
|
import cudf
|
|
|
|
nan_key_plan = build_gpu_aggregation_plan(
|
|
("item",), (Count(),), input_schema=pa.schema([("item", pa.null())])
|
|
)
|
|
assert isinstance(nan_key_plan, GPUAggregationPlan), nan_key_plan
|
|
normalized_nan_key_partial = nan_key_plan.partial_aggregate(
|
|
cudf.DataFrame({"item": cudf.Series([None], dtype="float64")}),
|
|
input_schema=pa.schema([("item", pa.null())]),
|
|
)
|
|
|
|
assert str(normalized_nan_key_partial["item"].dtype) == "float64"
|
|
assert (
|
|
str(normalized_nan_key_partial[nan_key_plan.accumulator_columns[0]].dtype)
|
|
== "int64"
|
|
)
|
|
|
|
def test_partial_aggregate_normalizes_unknown_schema_accumulators(
|
|
self, ray_with_gpu
|
|
):
|
|
import cudf
|
|
|
|
unknown_schema_plan = build_gpu_aggregation_plan(
|
|
("A",), (Sum("B"),), input_schema=pa.schema([("A", pa.int64())])
|
|
)
|
|
assert isinstance(unknown_schema_plan, GPUAggregationPlan), unknown_schema_plan
|
|
unknown_schema_acc_col = unknown_schema_plan.accumulator_columns[0]
|
|
int_input = cudf.DataFrame({"A": [0], "B": np.array([1], dtype=np.int64)})
|
|
double_input = cudf.DataFrame(
|
|
{"A": [0], "B": np.array([1.0], dtype=np.float64)}
|
|
)
|
|
|
|
normalized_int_partial = unknown_schema_plan.partial_aggregate(
|
|
int_input,
|
|
input_schema=pa.schema([("A", pa.int64()), ("B", pa.int64())]),
|
|
)
|
|
normalized_double_partial = unknown_schema_plan.partial_aggregate(
|
|
double_input,
|
|
input_schema=pa.schema([("A", pa.int64()), ("B", pa.float64())]),
|
|
)
|
|
|
|
assert str(normalized_int_partial[unknown_schema_acc_col].dtype) == "int64"
|
|
assert str(normalized_double_partial[unknown_schema_acc_col].dtype) == "float64"
|
|
|
|
def test_unsupported_cudf_dtype_cast_is_noop(self, monkeypatch):
|
|
cudf = pytest.importorskip("cudf")
|
|
|
|
df = cudf.DataFrame({"value": np.array([1, 2], dtype=np.int64)})
|
|
|
|
def raise_type_error(_self):
|
|
raise TypeError("unsupported cuDF dtype")
|
|
|
|
monkeypatch.setattr(hash_aggregate.DataType, "to_cudf_type", raise_type_error)
|
|
|
|
hash_aggregate._cast_cudf_column_dtype(
|
|
df, "value", hash_aggregate.DataType.from_numpy("int64")
|
|
)
|
|
|
|
assert df["value"].to_pandas().tolist() == [1, 2]
|
|
assert str(df["value"].dtype) == "int64"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPUHashAggregateActor — real GPU paths (conditional)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.gpu
|
|
class TestGPUHashAggregateActorReal:
|
|
"""Exercises GPU aggregate methods on actual hardware."""
|
|
|
|
def _make_setup_actor(self, aggregation_plan, total_nparts: int = 2):
|
|
actor_cls = cast(ActorClass[Any], GPUHashAggregateActor)
|
|
actor = cast(
|
|
ActorHandle[Any],
|
|
actor_cls.options(num_gpus=1).remote(
|
|
nranks=1,
|
|
total_nparts=total_nparts,
|
|
aggregation_plan=aggregation_plan,
|
|
),
|
|
)
|
|
|
|
_, root_address = ray.get(actor.setup_root.remote())
|
|
ray.get(actor.setup_worker.remote(root_address))
|
|
return actor
|
|
|
|
@staticmethod
|
|
def _collect_tables(actor) -> List[pa.Table]:
|
|
gen = actor.finish_and_extract.options(num_returns="streaming").remote()
|
|
return [
|
|
item for ref in gen for item in [ray.get(ref)] if isinstance(item, pa.Table)
|
|
]
|
|
|
|
@staticmethod
|
|
def _collect_frame(actor) -> pd.DataFrame:
|
|
tables = TestGPUHashAggregateActorReal._collect_tables(actor)
|
|
frames = [table.to_pandas() for table in tables if table.num_rows > 0]
|
|
if not frames:
|
|
return pd.DataFrame()
|
|
return pd.concat(frames, ignore_index=True)
|
|
|
|
def test_grouped_builtin_aggregates_real_gpu(self, ray_with_gpu):
|
|
table = pa.table(
|
|
{
|
|
"group": pa.array([0, 0, 1, 1, 1, 2], type=pa.int64()),
|
|
"value": pa.array([1, None, 2, 3, None, 10], type=pa.int64()),
|
|
}
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("group",),
|
|
(
|
|
Count(),
|
|
Count("value", ignore_nulls=True),
|
|
Sum("value"),
|
|
Min("value"),
|
|
Max("value"),
|
|
Mean("value"),
|
|
),
|
|
input_schema=table.schema,
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
actor = self._make_setup_actor(plan)
|
|
try:
|
|
assert ray.get(actor.insert_batch.remote(table)) == table.num_rows
|
|
result = (
|
|
self._collect_frame(actor).sort_values("group").reset_index(drop=True)
|
|
)
|
|
finally:
|
|
ray.kill(actor)
|
|
|
|
assert result["group"].tolist() == [0, 1, 2]
|
|
assert result["count()"].tolist() == [2, 3, 1]
|
|
assert result["count(value)"].tolist() == [1, 2, 1]
|
|
assert result["sum(value)"].tolist() == [1, 5, 10]
|
|
assert result["min(value)"].tolist() == [1, 2, 10]
|
|
assert result["max(value)"].tolist() == [1, 3, 10]
|
|
assert result["mean(value)"].tolist() == pytest.approx([1.0, 2.5, 10.0])
|
|
|
|
def test_grouped_aggregate_output_has_unique_join_keys(self, ray_with_gpu):
|
|
import cudf
|
|
|
|
schema = pa.schema(
|
|
[
|
|
("key1", pa.int64()),
|
|
("key2", pa.int64()),
|
|
("value", pa.int64()),
|
|
]
|
|
)
|
|
plan = build_gpu_aggregation_plan(
|
|
("key1", "key2"),
|
|
(
|
|
Min(on="value", alias_name="min_value"),
|
|
Count(alias_name="row_count"),
|
|
),
|
|
input_schema=schema,
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
df = cudf.DataFrame(
|
|
{
|
|
"key1": [0, 0, 1],
|
|
"key2": [10, 10, 20],
|
|
"value": [5, 6, 7],
|
|
}
|
|
)
|
|
partial = plan.partial_aggregate(df, input_schema=schema)
|
|
result = plan.final_aggregate(partial, input_schema=schema)
|
|
result_table = result.to_arrow(preserve_index=False)
|
|
|
|
assert result_table.column_names == ["key1", "key2", "min_value", "row_count"]
|
|
assert len(result_table.column_names) == len(set(result_table.column_names))
|
|
|
|
key_table = pa.table(
|
|
{
|
|
"key1": pa.array([0, 1], type=pa.int64()),
|
|
"key2": pa.array([10, 20], type=pa.int64()),
|
|
}
|
|
)
|
|
joined = key_table.join(
|
|
result_table,
|
|
keys=["key1", "key2"],
|
|
join_type="inner",
|
|
)
|
|
assert joined.num_rows == 2
|
|
|
|
def test_global_builtin_aggregates_real_gpu(self, ray_with_gpu):
|
|
plan = build_gpu_aggregation_plan(
|
|
tuple(),
|
|
(
|
|
Count(),
|
|
Count("value", ignore_nulls=True),
|
|
Sum("value"),
|
|
Mean("value"),
|
|
),
|
|
)
|
|
assert isinstance(plan, GPUAggregationPlan), plan
|
|
|
|
table = pa.table({"value": pa.array([1, None, 2, 5], type=pa.int64())})
|
|
actor = self._make_setup_actor(plan, total_nparts=1)
|
|
try:
|
|
assert ray.get(actor.insert_batch.remote(table)) == table.num_rows
|
|
result = self._collect_frame(actor)
|
|
finally:
|
|
ray.kill(actor)
|
|
|
|
assert result.columns.tolist() == [
|
|
"count()",
|
|
"count(value)",
|
|
"sum(value)",
|
|
"mean(value)",
|
|
]
|
|
assert len(result) == 1
|
|
row = result.iloc[0].to_dict()
|
|
assert row["count()"] == 4
|
|
assert row["count(value)"] == 3
|
|
assert row["sum(value)"] == 8
|
|
assert row["mean(value)"] == pytest.approx(8 / 3)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|