wehub-resource-sync
1691 lines
60 KiB
Python
1691 lines
60 KiB
Python
from __future__ import annotations
|
|
|
|
import abc
|
|
import logging
|
|
import pickle
|
|
import time
|
|
import types
|
|
import typing
|
|
from collections import defaultdict
|
|
from typing import (
|
|
Any,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Sequence,
|
|
Tuple,
|
|
Union,
|
|
cast,
|
|
)
|
|
|
|
import pyarrow as pa
|
|
|
|
import ray
|
|
from ray.data._internal.execution.interfaces import PhysicalOperator
|
|
from ray.data._internal.gpu_shuffle.hash_shuffle import (
|
|
_GPU_PARTITION_ID_KEY,
|
|
GPURankPool,
|
|
GPUShuffleOperator,
|
|
_derive_num_gpu_ranks,
|
|
)
|
|
from ray.data._internal.table_block import TableBlockAccessor
|
|
from ray.data.aggregate import AggregateFn, AggregateFnV2, Count, Max, Mean, Min, Sum
|
|
from ray.data.block import (
|
|
Block,
|
|
BlockAccessor,
|
|
BlockExecStats,
|
|
BlockMetadataWithSchema,
|
|
BlockStats,
|
|
Schema,
|
|
)
|
|
from ray.data.context import DataContext
|
|
from ray.data.datatype import DataType
|
|
|
|
if typing.TYPE_CHECKING:
|
|
import cudf
|
|
|
|
from ray.data._internal.progress.base_progress import BaseProgressBar
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
_GLOBAL_AGGREGATE_KEY = "__hash_aggregate_global_key"
|
|
|
|
|
|
def _cast_cudf_column_dtype(
|
|
df: cudf.DataFrame, column: str, dtype: Optional[DataType]
|
|
) -> None:
|
|
"""Cast a ``cudf.DataFrame`` column to specified dtype in place."""
|
|
if dtype is None or column not in df.columns:
|
|
return
|
|
dtype = DataType.from_dtype(dtype)
|
|
if dtype is None:
|
|
return
|
|
|
|
try:
|
|
cast_dtype = dtype.to_cudf_type()
|
|
except (TypeError, ValueError, NotImplementedError):
|
|
return
|
|
|
|
try:
|
|
df[column] = df[column].astype(cast_dtype)
|
|
except (TypeError, ValueError, NotImplementedError):
|
|
# fallback for handling all-null columns
|
|
if len(df) > 0 and not bool(df[column].isnull().all()):
|
|
return
|
|
|
|
import cudf
|
|
|
|
df[column] = cudf.Series([None] * len(df), dtype=cast_dtype)
|
|
|
|
|
|
def _cudf_column_dtype(df: cudf.DataFrame, column: str) -> Optional[DataType]:
|
|
"""Get the DataType of a column from a ``cudf.DataFrame``.
|
|
|
|
Returns None if the column is not found.
|
|
"""
|
|
if column not in df.columns:
|
|
return None
|
|
|
|
raw_dtype = df[column].dtype
|
|
try:
|
|
return DataType.from_cudf(raw_dtype)
|
|
except (AttributeError, ImportError, TypeError):
|
|
return DataType.from_dtype(raw_dtype)
|
|
|
|
|
|
def _schema_column_dtype(
|
|
schema: Optional[Schema], column: Optional[str]
|
|
) -> Optional[DataType]:
|
|
"""Get the DataType of a column from a ``Schema``.
|
|
|
|
Returns None if the column is not found (or None).
|
|
"""
|
|
if schema is None or column is None or column not in schema.names:
|
|
return None
|
|
|
|
if isinstance(schema, pa.Schema):
|
|
return DataType.from_dtype(schema.field(column).type)
|
|
|
|
return DataType.from_dtype(schema.types[schema.names.index(column)])
|
|
|
|
|
|
class GPUAggregateFn(abc.ABC):
|
|
"""Extension point for GPU-enabled aggregations.
|
|
|
|
GPU aggregate implementations define cuDF partial and final aggregation methods.
|
|
|
|
Args:
|
|
name: The name of the aggregation, which will be used as part of the column name
|
|
in the output, e.g. "sum" -> "sum(col)".
|
|
on: The name of the column to perform the aggregation on.
|
|
ignore_nulls: Whether to ignore null values during aggregation.
|
|
For example, should "count" include null rows or not?
|
|
accumulators: The names of (internal) accumulators used by the aggregation.
|
|
For example, "sum" uses "value" while "mean" uses ("sum", "count",
|
|
"null_count"). These will be combined with the accumulator_prefix to form
|
|
the final, unique names of the GPU accumulator columns.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
*,
|
|
on: Optional[str],
|
|
ignore_nulls: bool,
|
|
accumulators: Tuple[str, ...] = ("value",),
|
|
) -> None:
|
|
if not name:
|
|
raise ValueError(
|
|
f"Non-empty string has to be provided as name (got {name})"
|
|
)
|
|
if not accumulators or any(not accumulator for accumulator in accumulators):
|
|
raise ValueError("Accumulators must be non-empty strings.")
|
|
|
|
self.name = name
|
|
self.target_column = on
|
|
self.ignore_nulls = ignore_nulls
|
|
self._accumulators = accumulators
|
|
|
|
@abc.abstractmethod
|
|
def partial_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Any:
|
|
"""Aggregate one input block (as a ``cudf.DataFrame``) into GPU accumulator
|
|
columns."""
|
|
...
|
|
|
|
@abc.abstractmethod
|
|
def final_aggregate(
|
|
self,
|
|
df: Any,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
output_name: str,
|
|
) -> Any:
|
|
"""Aggregate shuffled GPU accumulator columns into final output."""
|
|
...
|
|
|
|
def _accumulator_columns(self, accumulator_prefix: str) -> Tuple[str, ...]:
|
|
"""Return the final, unique names of the GPU accumulator columns per
|
|
aggregation by concatenating the accumulator_prefix and the accumulator columns.
|
|
|
|
The accumulator_prefix is generated by the GPUAggregationPlan to uniquely
|
|
identify each GPU aggregation, e.g. for a single `mean` aggregation with
|
|
accumulator prefix "__ray_gpu_agg_0", the unique accumulator column names
|
|
will be:
|
|
- "__ray_gpu_agg_0_sum"
|
|
- "__ray_gpu_agg_0_count"
|
|
- "__ray_gpu_agg_0_null_count"
|
|
"""
|
|
return tuple(
|
|
f"{accumulator_prefix}_{accumulator}" for accumulator in self._accumulators
|
|
)
|
|
|
|
def _empty_global_partial_values(self, accumulator_prefix: str) -> Dict[str, Any]:
|
|
"""Return accumulator values for an empty block during a global aggregation
|
|
operation (no key columns).
|
|
|
|
This is used to ensure that the GPU aggregation can handle empty blocks
|
|
gracefully.
|
|
|
|
Subclasses should override this when all-null accumulator values are not
|
|
semantically correct for empty global input.
|
|
"""
|
|
return {
|
|
column: None for column in self._accumulator_columns(accumulator_prefix)
|
|
}
|
|
|
|
def _partial_accumulator_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
"""Return dtypes for partial accumulator columns.
|
|
|
|
Subclasses should override this when accumulator columns require explicit
|
|
dtype normalization.
|
|
"""
|
|
return {
|
|
column: None for column in self._accumulator_columns(accumulator_prefix)
|
|
}
|
|
|
|
def _final_cudf_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
output_name: str,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
"""Return dtypes for final cuDF output columns.
|
|
|
|
Subclasses should override this when final output columns require explicit
|
|
dtype normalization.
|
|
"""
|
|
return {}
|
|
|
|
def _final_arrow_types(
|
|
self,
|
|
output_name: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, pa.DataType]:
|
|
"""Return Arrow types for final output columns.
|
|
|
|
Subclasses should override this when Arrow output normalization requires
|
|
explicit types.
|
|
"""
|
|
return {}
|
|
|
|
|
|
def _fill_missing_count(
|
|
result: cudf.DataFrame, count_column: str, dtype: Optional[DataType] = None
|
|
) -> None:
|
|
"""Fill missing counts with 0 and cast to specified dtype.
|
|
|
|
This is used to ensure that the GPU aggregation can handle empty blocks
|
|
gracefully.
|
|
"""
|
|
if count_column not in result.columns:
|
|
result[count_column] = 0
|
|
else:
|
|
result[count_column] = result[count_column].fillna(0)
|
|
_cast_cudf_column_dtype(result, count_column, dtype)
|
|
|
|
|
|
def _fill_missing_reduction(
|
|
result: cudf.DataFrame, reduction_column: str, dtype: Optional[DataType] = None
|
|
) -> None:
|
|
"""Fill any missing reduction values with None and cast to specified dtype.
|
|
|
|
This is used to ensure that the GPU aggregation can handle empty blocks
|
|
gracefully.
|
|
"""
|
|
if reduction_column not in result.columns:
|
|
result[reduction_column] = None
|
|
_cast_cudf_column_dtype(result, reduction_column, dtype)
|
|
|
|
|
|
class GPUCount(GPUAggregateFn):
|
|
"""GPU implementation for :class:`ray.data.aggregate.Count`."""
|
|
|
|
def __init__(self, agg: Count, *, source_dtype: Optional[DataType] = None) -> None:
|
|
self.source_dtype = source_dtype
|
|
super().__init__(
|
|
agg.name,
|
|
on=agg.get_target_column(),
|
|
ignore_nulls=agg._ignore_nulls,
|
|
accumulators=("value",),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> cudf.DataFrame:
|
|
acc_col = accumulator_columns[0]
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
if self.target_column is None or not self.ignore_nulls:
|
|
result = grouped.size().reset_index()
|
|
return result.rename(columns={result.columns[-1]: acc_col})
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: acc_col})
|
|
count_dtype = _cudf_column_dtype(sizes, acc_col)
|
|
|
|
counts = grouped[self.target_column].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: acc_col})
|
|
|
|
result = sizes[list(key_columns)].merge(
|
|
counts, on=list(key_columns), how="left"
|
|
)
|
|
_fill_missing_count(result, acc_col, count_dtype)
|
|
return result[list(key_columns) + [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]
|
|
result = (
|
|
df.groupby(list(key_columns), dropna=False)[acc_col].sum().reset_index()
|
|
)
|
|
result = result.rename(columns={result.columns[-1]: output_name})
|
|
return result[list(key_columns) + [output_name]]
|
|
|
|
def _empty_global_partial_values(self, accumulator_prefix: str) -> Dict[str, Any]:
|
|
return {self._accumulator_columns(accumulator_prefix)[0]: 0}
|
|
|
|
def _partial_accumulator_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
return {
|
|
self._accumulator_columns(accumulator_prefix)[0]: DataType.from_numpy(
|
|
"int64"
|
|
)
|
|
}
|
|
|
|
def _final_cudf_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
output_name: str,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
return {output_name: DataType.from_numpy("int64")}
|
|
|
|
|
|
class GPUSum(GPUAggregateFn):
|
|
"""GPU implementation for :class:`ray.data.aggregate.Sum`."""
|
|
|
|
def __init__(self, agg: Sum, *, source_dtype: Optional[DataType] = None) -> None:
|
|
self.source_dtype = source_dtype
|
|
super().__init__(
|
|
agg.name,
|
|
on=agg.get_target_column(),
|
|
ignore_nulls=agg._ignore_nulls,
|
|
accumulators=("value",),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> cudf.DataFrame:
|
|
assert self.target_column is not None
|
|
acc_col = accumulator_columns[0]
|
|
output_dtype = self._reduction_dtype(df, input_schema)
|
|
size_col = f"{acc_col}_size"
|
|
count_col = f"{acc_col}_count"
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: size_col})
|
|
count_dtype = _cudf_column_dtype(sizes, size_col)
|
|
|
|
counts = grouped[self.target_column].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: count_col})
|
|
|
|
result = sizes.merge(counts, on=list(key_columns), how="left")
|
|
_fill_missing_count(result, count_col, count_dtype)
|
|
|
|
if len(result) > 0 and bool(cast("cudf.Series", result[count_col] == 0).all()):
|
|
result[acc_col] = None
|
|
_cast_cudf_column_dtype(result, acc_col, output_dtype)
|
|
else:
|
|
aggregated = grouped[self.target_column].sum().reset_index()
|
|
aggregated = aggregated.rename(columns={aggregated.columns[-1]: acc_col})
|
|
result = result.merge(aggregated, on=list(key_columns), how="left")
|
|
_fill_missing_reduction(result, acc_col, output_dtype)
|
|
|
|
null_mask = result[count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = result[size_col] != result[count_col]
|
|
result.loc[null_mask, acc_col] = None
|
|
return result[list(key_columns) + [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]
|
|
output_dtype = _cudf_column_dtype(df, acc_col)
|
|
size_col = f"{acc_col}_partial_size"
|
|
count_col = f"{acc_col}_partial_count"
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: size_col})
|
|
count_dtype = _cudf_column_dtype(sizes, size_col)
|
|
|
|
counts = grouped[acc_col].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: count_col})
|
|
|
|
result = sizes.merge(counts, on=list(key_columns), how="left")
|
|
_fill_missing_count(result, count_col, count_dtype)
|
|
|
|
if len(result) > 0 and bool(cast("cudf.Series", result[count_col] == 0).all()):
|
|
result[output_name] = None
|
|
_cast_cudf_column_dtype(result, output_name, output_dtype)
|
|
else:
|
|
aggregated = grouped[acc_col].sum().reset_index()
|
|
aggregated = aggregated.rename(
|
|
columns={aggregated.columns[-1]: output_name}
|
|
)
|
|
result = result.merge(aggregated, on=list(key_columns), how="left")
|
|
_fill_missing_reduction(result, output_name, output_dtype)
|
|
|
|
null_mask = result[count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = result[size_col] != result[count_col]
|
|
result.loc[null_mask, output_name] = None
|
|
return result[list(key_columns) + [output_name]]
|
|
|
|
def _empty_global_partial_values(self, accumulator_prefix: str) -> Dict[str, Any]:
|
|
return {
|
|
self._accumulator_columns(accumulator_prefix)[0]: (
|
|
None if self.ignore_nulls else 0
|
|
)
|
|
}
|
|
|
|
def _partial_accumulator_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
assert self.target_column is not None
|
|
return {
|
|
self._accumulator_columns(accumulator_prefix)[0]: self._reduction_dtype(
|
|
df, input_schema
|
|
)
|
|
}
|
|
|
|
def _reduction_dtype(
|
|
self, df: cudf.DataFrame, input_schema: Optional[Schema]
|
|
) -> Optional[DataType]:
|
|
assert self.target_column is not None
|
|
dtype = self.source_dtype
|
|
if dtype is None:
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None:
|
|
dtype = _cudf_column_dtype(df, self.target_column)
|
|
if dtype is None:
|
|
return None
|
|
if dtype.is_null_type() or dtype.is_boolean_type():
|
|
return DataType.from_numpy("int64")
|
|
if dtype.is_integer_type():
|
|
return DataType.from_numpy("uint64" if dtype.is_uint64_type() else "int64")
|
|
if dtype.is_floating_type():
|
|
return DataType.from_numpy("float64")
|
|
return dtype
|
|
|
|
def _final_cudf_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
output_name: str,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
acc_col = self._accumulator_columns(accumulator_prefix)[0]
|
|
acc_dtype = _cudf_column_dtype(df, acc_col)
|
|
if acc_dtype is None:
|
|
acc_dtype = self._partial_accumulator_dtypes(
|
|
df, accumulator_prefix, input_schema=input_schema
|
|
)[acc_col]
|
|
return {output_name: acc_dtype}
|
|
|
|
def _final_arrow_types(
|
|
self,
|
|
output_name: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, pa.DataType]:
|
|
dtype = self.source_dtype
|
|
if dtype is None or dtype.is_null_type():
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None or not dtype.is_null_type():
|
|
return {}
|
|
return {output_name: pa.null()}
|
|
|
|
|
|
class GPUMin(GPUAggregateFn):
|
|
"""GPU implementation for :class:`ray.data.aggregate.Min`."""
|
|
|
|
def __init__(self, agg: Min, *, source_dtype: Optional[DataType] = None) -> None:
|
|
self.source_dtype = source_dtype
|
|
super().__init__(
|
|
agg.name,
|
|
on=agg.get_target_column(),
|
|
ignore_nulls=agg._ignore_nulls,
|
|
accumulators=("value",),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> cudf.DataFrame:
|
|
assert self.target_column is not None
|
|
acc_col = accumulator_columns[0]
|
|
output_dtype = self._reduction_dtype(df, input_schema)
|
|
size_col = f"{acc_col}_size"
|
|
count_col = f"{acc_col}_count"
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: size_col})
|
|
count_dtype = _cudf_column_dtype(sizes, size_col)
|
|
|
|
counts = grouped[self.target_column].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: count_col})
|
|
|
|
result = sizes.merge(counts, on=list(key_columns), how="left")
|
|
_fill_missing_count(result, count_col, count_dtype)
|
|
|
|
if len(result) > 0 and bool(cast("cudf.Series", result[count_col] == 0).all()):
|
|
result[acc_col] = None
|
|
_cast_cudf_column_dtype(result, acc_col, output_dtype)
|
|
else:
|
|
aggregated = grouped[self.target_column].min().reset_index()
|
|
aggregated = aggregated.rename(columns={aggregated.columns[-1]: acc_col})
|
|
result = result.merge(aggregated, on=list(key_columns), how="left")
|
|
_fill_missing_reduction(result, acc_col, output_dtype)
|
|
|
|
null_mask = result[count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = result[size_col] != result[count_col]
|
|
result.loc[null_mask, acc_col] = None
|
|
return result[list(key_columns) + [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]
|
|
output_dtype = _cudf_column_dtype(df, acc_col)
|
|
size_col = f"{acc_col}_partial_size"
|
|
count_col = f"{acc_col}_partial_count"
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: size_col})
|
|
count_dtype = _cudf_column_dtype(sizes, size_col)
|
|
|
|
counts = grouped[acc_col].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: count_col})
|
|
|
|
result = sizes.merge(counts, on=list(key_columns), how="left")
|
|
_fill_missing_count(result, count_col, count_dtype)
|
|
|
|
if len(result) > 0 and bool(cast("cudf.Series", result[count_col] == 0).all()):
|
|
result[output_name] = None
|
|
_cast_cudf_column_dtype(result, output_name, output_dtype)
|
|
else:
|
|
aggregated = grouped[acc_col].min().reset_index()
|
|
aggregated = aggregated.rename(
|
|
columns={aggregated.columns[-1]: output_name}
|
|
)
|
|
result = result.merge(aggregated, on=list(key_columns), how="left")
|
|
_fill_missing_reduction(result, output_name, output_dtype)
|
|
|
|
null_mask = result[count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = result[size_col] != result[count_col]
|
|
result.loc[null_mask, output_name] = None
|
|
return result[list(key_columns) + [output_name]]
|
|
|
|
def _empty_global_partial_values(self, accumulator_prefix: str) -> Dict[str, Any]:
|
|
return {
|
|
self._accumulator_columns(accumulator_prefix)[0]: (
|
|
None if self.ignore_nulls else float("+inf")
|
|
)
|
|
}
|
|
|
|
def _partial_accumulator_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
assert self.target_column is not None
|
|
return {
|
|
self._accumulator_columns(accumulator_prefix)[0]: self._reduction_dtype(
|
|
df, input_schema
|
|
)
|
|
}
|
|
|
|
def _reduction_dtype(
|
|
self, df: cudf.DataFrame, input_schema: Optional[Schema]
|
|
) -> Optional[DataType]:
|
|
assert self.target_column is not None
|
|
dtype = self.source_dtype
|
|
if dtype is None:
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None:
|
|
dtype = _cudf_column_dtype(df, self.target_column)
|
|
if dtype is None:
|
|
return None
|
|
if dtype.is_null_type():
|
|
return DataType.from_numpy("int64")
|
|
if dtype.is_boolean_type():
|
|
return DataType.from_numpy("bool")
|
|
if dtype.is_integer_type():
|
|
return DataType.from_numpy("uint64" if dtype.is_uint64_type() else "int64")
|
|
if dtype.is_floating_type():
|
|
return DataType.from_numpy("float64")
|
|
return dtype
|
|
|
|
def _final_cudf_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
output_name: str,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
acc_col = self._accumulator_columns(accumulator_prefix)[0]
|
|
acc_dtype = _cudf_column_dtype(df, acc_col)
|
|
if acc_dtype is None:
|
|
acc_dtype = self._partial_accumulator_dtypes(
|
|
df, accumulator_prefix, input_schema=input_schema
|
|
)[acc_col]
|
|
return {output_name: acc_dtype}
|
|
|
|
def _final_arrow_types(
|
|
self,
|
|
output_name: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, pa.DataType]:
|
|
dtype = self.source_dtype
|
|
if dtype is None or dtype.is_null_type():
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None or not dtype.is_null_type():
|
|
return {}
|
|
return {output_name: pa.null()}
|
|
|
|
|
|
class GPUMax(GPUAggregateFn):
|
|
"""GPU implementation for :class:`ray.data.aggregate.Max`."""
|
|
|
|
def __init__(self, agg: Max, *, source_dtype: Optional[DataType] = None) -> None:
|
|
self.source_dtype = source_dtype
|
|
super().__init__(
|
|
agg.name,
|
|
on=agg.get_target_column(),
|
|
ignore_nulls=agg._ignore_nulls,
|
|
accumulators=("value",),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> cudf.DataFrame:
|
|
assert self.target_column is not None
|
|
acc_col = accumulator_columns[0]
|
|
output_dtype = self._reduction_dtype(df, input_schema)
|
|
size_col = f"{acc_col}_size"
|
|
count_col = f"{acc_col}_count"
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: size_col})
|
|
count_dtype = _cudf_column_dtype(sizes, size_col)
|
|
|
|
counts = grouped[self.target_column].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: count_col})
|
|
|
|
result = sizes.merge(counts, on=list(key_columns), how="left")
|
|
_fill_missing_count(result, count_col, count_dtype)
|
|
|
|
if len(result) > 0 and bool(cast("cudf.Series", result[count_col] == 0).all()):
|
|
result[acc_col] = None
|
|
_cast_cudf_column_dtype(result, acc_col, output_dtype)
|
|
else:
|
|
aggregated = grouped[self.target_column].max().reset_index()
|
|
aggregated = aggregated.rename(columns={aggregated.columns[-1]: acc_col})
|
|
result = result.merge(aggregated, on=list(key_columns), how="left")
|
|
_fill_missing_reduction(result, acc_col, output_dtype)
|
|
|
|
null_mask = result[count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = result[size_col] != result[count_col]
|
|
result.loc[null_mask, acc_col] = None
|
|
return result[list(key_columns) + [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]
|
|
output_dtype = _cudf_column_dtype(df, acc_col)
|
|
size_col = f"{acc_col}_partial_size"
|
|
count_col = f"{acc_col}_partial_count"
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: size_col})
|
|
count_dtype = _cudf_column_dtype(sizes, size_col)
|
|
|
|
counts = grouped[acc_col].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: count_col})
|
|
|
|
result = sizes.merge(counts, on=list(key_columns), how="left")
|
|
_fill_missing_count(result, count_col, count_dtype)
|
|
|
|
if len(result) > 0 and bool(cast("cudf.Series", result[count_col] == 0).all()):
|
|
result[output_name] = None
|
|
_cast_cudf_column_dtype(result, output_name, output_dtype)
|
|
else:
|
|
aggregated = grouped[acc_col].max().reset_index()
|
|
aggregated = aggregated.rename(
|
|
columns={aggregated.columns[-1]: output_name}
|
|
)
|
|
result = result.merge(aggregated, on=list(key_columns), how="left")
|
|
_fill_missing_reduction(result, output_name, output_dtype)
|
|
|
|
null_mask = result[count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = result[size_col] != result[count_col]
|
|
result.loc[null_mask, output_name] = None
|
|
return result[list(key_columns) + [output_name]]
|
|
|
|
def _empty_global_partial_values(self, accumulator_prefix: str) -> Dict[str, Any]:
|
|
return {
|
|
self._accumulator_columns(accumulator_prefix)[0]: (
|
|
None if self.ignore_nulls else float("-inf")
|
|
)
|
|
}
|
|
|
|
def _partial_accumulator_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
assert self.target_column is not None
|
|
return {
|
|
self._accumulator_columns(accumulator_prefix)[0]: self._reduction_dtype(
|
|
df, input_schema
|
|
)
|
|
}
|
|
|
|
def _reduction_dtype(
|
|
self, df: cudf.DataFrame, input_schema: Optional[Schema]
|
|
) -> Optional[DataType]:
|
|
assert self.target_column is not None
|
|
dtype = self.source_dtype
|
|
if dtype is None:
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None:
|
|
dtype = _cudf_column_dtype(df, self.target_column)
|
|
if dtype is None:
|
|
return None
|
|
if dtype.is_null_type():
|
|
return DataType.from_numpy("int64")
|
|
if dtype.is_boolean_type():
|
|
return DataType.from_numpy("bool")
|
|
if dtype.is_integer_type():
|
|
return DataType.from_numpy("uint64" if dtype.is_uint64_type() else "int64")
|
|
if dtype.is_floating_type():
|
|
return DataType.from_numpy("float64")
|
|
return dtype
|
|
|
|
def _final_cudf_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
output_name: str,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
acc_col = self._accumulator_columns(accumulator_prefix)[0]
|
|
acc_dtype = _cudf_column_dtype(df, acc_col)
|
|
if acc_dtype is None:
|
|
acc_dtype = self._partial_accumulator_dtypes(
|
|
df, accumulator_prefix, input_schema=input_schema
|
|
)[acc_col]
|
|
return {output_name: acc_dtype}
|
|
|
|
def _final_arrow_types(
|
|
self,
|
|
output_name: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, pa.DataType]:
|
|
dtype = self.source_dtype
|
|
if dtype is None or dtype.is_null_type():
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None or not dtype.is_null_type():
|
|
return {}
|
|
return {output_name: pa.null()}
|
|
|
|
|
|
class GPUMean(GPUAggregateFn):
|
|
"""GPU implementation for :class:`ray.data.aggregate.Mean`."""
|
|
|
|
def __init__(self, agg: Mean, *, source_dtype: Optional[DataType] = None) -> None:
|
|
self.source_dtype = source_dtype
|
|
super().__init__(
|
|
agg.name,
|
|
on=agg.get_target_column(),
|
|
ignore_nulls=agg._ignore_nulls,
|
|
accumulators=("sum", "count", "null_count"),
|
|
)
|
|
|
|
def partial_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> cudf.DataFrame:
|
|
assert self.target_column is not None
|
|
|
|
sum_col, count_col, null_count_col = accumulator_columns
|
|
size_col = f"{sum_col}_size"
|
|
output_dtype = self._reduction_dtype(df, input_schema)
|
|
grouped = df.groupby(list(key_columns), dropna=False)
|
|
|
|
sizes = grouped.size().reset_index()
|
|
sizes = sizes.rename(columns={sizes.columns[-1]: size_col})
|
|
count_dtype = _cudf_column_dtype(sizes, size_col)
|
|
|
|
counts = grouped[self.target_column].count().reset_index()
|
|
counts = counts.rename(columns={counts.columns[-1]: count_col})
|
|
|
|
result = sizes.merge(counts, on=list(key_columns), how="left")
|
|
_fill_missing_count(result, count_col, count_dtype)
|
|
|
|
if len(result) > 0 and bool(cast("cudf.Series", result[count_col] == 0).all()):
|
|
result[sum_col] = None
|
|
_cast_cudf_column_dtype(result, sum_col, output_dtype)
|
|
else:
|
|
aggregated = grouped[self.target_column].sum().reset_index()
|
|
aggregated = aggregated.rename(columns={aggregated.columns[-1]: sum_col})
|
|
result = result.merge(aggregated, on=list(key_columns), how="left")
|
|
_fill_missing_reduction(result, sum_col, output_dtype)
|
|
|
|
null_mask = result[count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = result[size_col] != result[count_col]
|
|
result.loc[null_mask, sum_col] = None
|
|
|
|
result[null_count_col] = result[size_col] - result[count_col]
|
|
_cast_cudf_column_dtype(result, null_count_col, count_dtype)
|
|
|
|
return result[list(key_columns) + list(accumulator_columns)]
|
|
|
|
def final_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
accumulator_columns: Tuple[str, ...],
|
|
output_name: str,
|
|
) -> cudf.DataFrame:
|
|
sum_col, count_col, null_count_col = accumulator_columns
|
|
final_sum_col = f"{sum_col}_final_sum"
|
|
final_count_col = f"{count_col}_final_count"
|
|
final_null_count_col = f"{null_count_col}_final_null_count"
|
|
sum_dtype = _cudf_column_dtype(df, sum_col)
|
|
|
|
acc_cols = [count_col, null_count_col, sum_col]
|
|
aggregated = (
|
|
df.groupby(list(key_columns), dropna=False)[acc_cols].sum().reset_index()
|
|
)
|
|
result = aggregated.rename(
|
|
columns={
|
|
count_col: final_count_col,
|
|
null_count_col: final_null_count_col,
|
|
sum_col: final_sum_col,
|
|
}
|
|
)
|
|
_fill_missing_reduction(result, final_sum_col, sum_dtype)
|
|
|
|
result[output_name] = result[final_sum_col] / result[final_count_col]
|
|
|
|
null_mask = result[final_count_col] == 0
|
|
if not self.ignore_nulls:
|
|
null_mask = null_mask | (result[final_null_count_col] > 0)
|
|
result.loc[null_mask, output_name] = None
|
|
|
|
return result[list(key_columns) + [output_name]]
|
|
|
|
def _empty_global_partial_values(self, accumulator_prefix: str) -> Dict[str, Any]:
|
|
sum_col, count_col, null_count_col = self._accumulator_columns(
|
|
accumulator_prefix
|
|
)
|
|
return {sum_col: None, count_col: 0, null_count_col: 0}
|
|
|
|
def _partial_accumulator_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
sum_col, count_col, null_count_col = self._accumulator_columns(
|
|
accumulator_prefix
|
|
)
|
|
assert self.target_column is not None
|
|
return {
|
|
sum_col: self._reduction_dtype(df, input_schema),
|
|
count_col: DataType.from_numpy("int64"),
|
|
null_count_col: DataType.from_numpy("int64"),
|
|
}
|
|
|
|
def _reduction_dtype(
|
|
self, df: cudf.DataFrame, input_schema: Optional[Schema]
|
|
) -> Optional[DataType]:
|
|
assert self.target_column is not None
|
|
dtype = self.source_dtype
|
|
if dtype is None:
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None:
|
|
dtype = _cudf_column_dtype(df, self.target_column)
|
|
if dtype is None:
|
|
return None
|
|
if dtype.is_null_type():
|
|
return DataType.from_numpy("float64")
|
|
if dtype.is_boolean_type():
|
|
return DataType.from_numpy("int64")
|
|
if dtype.is_integer_type():
|
|
return DataType.from_numpy("uint64" if dtype.is_uint64_type() else "int64")
|
|
if dtype.is_floating_type():
|
|
return DataType.from_numpy("float64")
|
|
return dtype
|
|
|
|
def _final_cudf_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
output_name: str,
|
|
accumulator_prefix: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
return {output_name: DataType.from_numpy("float64")}
|
|
|
|
def _final_arrow_types(
|
|
self,
|
|
output_name: str,
|
|
*,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, pa.DataType]:
|
|
dtype = self.source_dtype
|
|
if dtype is None or dtype.is_null_type():
|
|
dtype = _schema_column_dtype(input_schema, self.target_column)
|
|
if dtype is None or not dtype.is_null_type():
|
|
return {}
|
|
return {output_name: pa.null()}
|
|
|
|
|
|
def _empty_dataframe(
|
|
cudf_module: types.ModuleType,
|
|
columns: Sequence[str],
|
|
dtypes: Optional[Dict[str, Optional[DataType]]] = None,
|
|
) -> cudf.DataFrame:
|
|
"""Create an empty ``cudf.DataFrame`` with specified columns and dtypes."""
|
|
dtypes = dtypes or {}
|
|
df = cudf_module.DataFrame()
|
|
for column in columns:
|
|
df[column] = []
|
|
_cast_cudf_column_dtype(df, column, dtypes.get(column))
|
|
return df
|
|
|
|
|
|
class GPUAggregationPlan:
|
|
"""Executable GPU aggregation plan shared by the driver and GPU actors.
|
|
|
|
Args:
|
|
key_columns: The key columns to group by.
|
|
gpu_aggregates: The GPU aggregate functions to apply.
|
|
accumulator_prefix: The prefix for intermediate accumulator columns.
|
|
input_schema: The schema of the input data.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
key_columns: Tuple[str, ...],
|
|
gpu_aggregates: Tuple[GPUAggregateFn, ...],
|
|
accumulator_prefix: str,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> None:
|
|
if not accumulator_prefix:
|
|
raise ValueError("Accumulator prefix must be a non-empty string.")
|
|
|
|
self._key_columns = key_columns
|
|
self._gpu_aggregates = gpu_aggregates
|
|
self._input_schema = input_schema
|
|
self._is_global = not key_columns
|
|
self._shuffle_key_columns = key_columns
|
|
|
|
# Resolve duplicate aggregation names the same way TableBlockAccessor does.
|
|
counts: Dict[str, int] = defaultdict(int)
|
|
resolved_names: List[str] = []
|
|
for agg in gpu_aggregates:
|
|
name = agg.name
|
|
if counts[name] > 0:
|
|
name = TableBlockAccessor._munge_conflict(name, counts[name])
|
|
counts[agg.name] += 1
|
|
resolved_names.append(name)
|
|
self._output_names = tuple(resolved_names)
|
|
|
|
# Generate unique accumulator prefixes for each resolved name
|
|
self._accumulator_prefixes = tuple(
|
|
f"{accumulator_prefix}_{index}" for index, _ in enumerate(gpu_aggregates)
|
|
)
|
|
|
|
# If global aggregation (w/o key columns), use an artificial shuffle key.
|
|
if self._is_global:
|
|
# filter out empty target columns
|
|
required_columns = {
|
|
agg.target_column
|
|
for agg in gpu_aggregates
|
|
if agg.target_column is not None
|
|
}
|
|
# ensure a unique global key by prepending an underscore if needed
|
|
# (just in case there is a collision)
|
|
global_key = _GLOBAL_AGGREGATE_KEY
|
|
while global_key in required_columns:
|
|
global_key = f"_{global_key}"
|
|
# set the shuffle key to the global key
|
|
self._shuffle_key_columns = (global_key,)
|
|
|
|
@property
|
|
def accumulator_columns(self) -> Tuple[str, ...]:
|
|
"""Return all internal accumulator columns for the GPU aggregation plan."""
|
|
columns: List[str] = []
|
|
for agg, accumulator_prefix in zip(
|
|
self._gpu_aggregates, self._accumulator_prefixes
|
|
):
|
|
columns.extend(agg._accumulator_columns(accumulator_prefix))
|
|
return tuple(columns)
|
|
|
|
@property
|
|
def output_names(self) -> Tuple[str, ...]:
|
|
"""Return all final output names for the GPU aggregation plan.
|
|
|
|
These will be used as the column names for the final output of the GPU
|
|
aggregations in the plan, e.g. "sum(col1)", "mean(col2)", etc.
|
|
"""
|
|
return self._output_names
|
|
|
|
@property
|
|
def required_columns(self) -> Tuple[str, ...]:
|
|
"""Return all required columns for the GPU aggregation plan.
|
|
|
|
These include the key columns and aggregation target columns, e.g.
|
|
groupby("col1").sum("col2")
|
|
"""
|
|
columns = list(self._key_columns)
|
|
for agg in self._gpu_aggregates:
|
|
target_column = agg.target_column
|
|
if target_column is not None and target_column not in columns:
|
|
columns.append(target_column)
|
|
return tuple(columns)
|
|
|
|
@property
|
|
def shuffle_key_columns(self) -> Tuple[str, ...]:
|
|
"""Return the shuffle key columns for the GPU aggregation plan."""
|
|
return self._shuffle_key_columns
|
|
|
|
def normalize_output_arrow(
|
|
self,
|
|
table: pa.Table,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> pa.Table:
|
|
arrow_types = self._final_arrow_types(input_schema)
|
|
if not arrow_types:
|
|
return table
|
|
|
|
columns = []
|
|
for column_name in table.column_names:
|
|
if column_name in arrow_types:
|
|
columns.append(pa.nulls(table.num_rows, type=arrow_types[column_name]))
|
|
else:
|
|
columns.append(table[column_name])
|
|
return pa.table(columns, names=table.column_names)
|
|
|
|
def partial_aggregate(
|
|
self, df: cudf.DataFrame, input_schema: Optional[Schema] = None
|
|
) -> cudf.DataFrame:
|
|
import cudf as cudf_module
|
|
|
|
if self._is_global:
|
|
df = df.copy(deep=False)
|
|
df[self._shuffle_key_columns[0]] = 0
|
|
|
|
key_columns = self._shuffle_key_columns
|
|
if len(df) == 0:
|
|
if self._is_global:
|
|
values: Dict[str, List[Any]] = {key_columns[0]: [0]}
|
|
for agg, accumulator_prefix in zip(
|
|
self._gpu_aggregates,
|
|
self._accumulator_prefixes,
|
|
):
|
|
empty_values = agg._empty_global_partial_values(accumulator_prefix)
|
|
for column, value in empty_values.items():
|
|
values[column] = [value]
|
|
result = cudf_module.DataFrame(values)[
|
|
list(key_columns) + list(self.accumulator_columns)
|
|
]
|
|
for column, dtype in self._partial_accumulator_dtypes(
|
|
df, key_columns, input_schema=input_schema
|
|
).items():
|
|
_cast_cudf_column_dtype(result, column, dtype)
|
|
return result
|
|
return _empty_dataframe(
|
|
cudf_module,
|
|
list(key_columns) + list(self.accumulator_columns),
|
|
dtypes=self._partial_accumulator_dtypes(
|
|
df, key_columns, input_schema=input_schema
|
|
),
|
|
)
|
|
|
|
result = None
|
|
for agg, accumulator_prefix in zip(
|
|
self._gpu_aggregates, self._accumulator_prefixes
|
|
):
|
|
accumulator_columns = agg._accumulator_columns(accumulator_prefix)
|
|
partial = agg.partial_aggregate(
|
|
df,
|
|
key_columns,
|
|
accumulator_columns,
|
|
input_schema=input_schema,
|
|
)
|
|
result = (
|
|
partial
|
|
if result is None
|
|
else result.merge(partial, on=list(key_columns), how="outer")
|
|
)
|
|
|
|
assert result is not None
|
|
for column, dtype in self._partial_accumulator_dtypes(
|
|
df, key_columns, input_schema=input_schema
|
|
).items():
|
|
_cast_cudf_column_dtype(result, column, dtype)
|
|
return result[list(key_columns) + list(self.accumulator_columns)]
|
|
|
|
def final_aggregate(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> cudf.DataFrame:
|
|
import cudf as cudf_module
|
|
|
|
key_columns = self._shuffle_key_columns
|
|
output_columns = ([] if self._is_global else list(key_columns)) + list(
|
|
self.output_names
|
|
)
|
|
|
|
if len(df) == 0:
|
|
return _empty_dataframe(
|
|
cudf_module,
|
|
output_columns,
|
|
dtypes=self._final_cudf_dtypes(
|
|
df,
|
|
input_schema=input_schema,
|
|
),
|
|
)
|
|
|
|
result = None
|
|
for agg, output_name, accumulator_prefix in zip(
|
|
self._gpu_aggregates, self._output_names, self._accumulator_prefixes
|
|
):
|
|
accumulator_columns = agg._accumulator_columns(accumulator_prefix)
|
|
finalized = agg.final_aggregate(
|
|
df,
|
|
key_columns,
|
|
accumulator_columns,
|
|
output_name,
|
|
)
|
|
result = (
|
|
finalized
|
|
if result is None
|
|
else result.merge(finalized, on=list(key_columns), how="outer")
|
|
)
|
|
|
|
assert result is not None
|
|
if self._is_global:
|
|
result = result.drop(columns=[self._shuffle_key_columns[0]])
|
|
|
|
return result[output_columns]
|
|
|
|
def merge_input_schema(
|
|
self, current: Optional[pa.Schema], observed: Optional[Schema]
|
|
) -> Optional[pa.Schema]:
|
|
"""Merge an observed block schema into the current runtime input schema."""
|
|
if observed is None:
|
|
return current
|
|
|
|
fields: Dict[str, pa.DataType] = {}
|
|
if current is not None:
|
|
fields.update({field.name: field.type for field in current})
|
|
|
|
for column in self.required_columns:
|
|
observed_dtype = _schema_column_dtype(observed, column)
|
|
if observed_dtype is None:
|
|
continue
|
|
|
|
try:
|
|
observed_arrow_dtype = observed_dtype.to_arrow_dtype()
|
|
except (AssertionError, TypeError, pa.ArrowNotImplementedError):
|
|
continue
|
|
|
|
current_dtype = fields.get(column)
|
|
if current_dtype is None or pa.types.is_null(current_dtype):
|
|
fields[column] = observed_arrow_dtype
|
|
elif not current_dtype.equals(observed_arrow_dtype):
|
|
# Unify schemas using arrow_ops
|
|
try:
|
|
from ray.data._internal.arrow_ops.transform_pyarrow import (
|
|
unify_schemas,
|
|
)
|
|
|
|
fields[column] = (
|
|
unify_schemas(
|
|
[
|
|
pa.schema([(column, current_dtype)]),
|
|
pa.schema([(column, observed_arrow_dtype)]),
|
|
],
|
|
promote_types=True,
|
|
)
|
|
.field(column)
|
|
.type
|
|
)
|
|
except (pa.ArrowInvalid, pa.ArrowTypeError):
|
|
pass
|
|
|
|
if not fields:
|
|
return current
|
|
|
|
ordered_names = [column for column in self.required_columns if column in fields]
|
|
ordered_names.extend(name for name in fields if name not in ordered_names)
|
|
return pa.schema([(name, fields[name]) for name in ordered_names])
|
|
|
|
def _effective_column_dtype(
|
|
self, column: str, runtime_input_schema: Optional[Schema] = None
|
|
) -> Optional[DataType]:
|
|
"""Return the dtype for a column in the input schema.
|
|
|
|
This method first checks the input schema provided to the GPUAggregationPlan
|
|
constructor, and then falls back to the runtime input schema if provided.
|
|
"""
|
|
dtype = _schema_column_dtype(self._input_schema, column)
|
|
if dtype is not None:
|
|
return dtype
|
|
return _schema_column_dtype(runtime_input_schema, column)
|
|
|
|
def _effective_input_schema(
|
|
self, runtime_input_schema: Optional[Schema] = None
|
|
) -> Optional[Schema]:
|
|
"""Return the effective input schema for the GPU aggregation plan.
|
|
|
|
This method supplies a fallback schema derived from the input schema,
|
|
if the runtime input schema is not provided.
|
|
"""
|
|
return (
|
|
runtime_input_schema
|
|
if runtime_input_schema is not None
|
|
else self._input_schema
|
|
)
|
|
|
|
def _final_arrow_types(
|
|
self, input_schema: Optional[Schema] = None
|
|
) -> Dict[str, pa.DataType]:
|
|
"""Return the Arrow types for the final output columns of the GPU aggregation plan."""
|
|
input_schema = self._effective_input_schema(input_schema)
|
|
types: Dict[str, pa.DataType] = {}
|
|
for agg, output_name in zip(self._gpu_aggregates, self._output_names):
|
|
types.update(
|
|
agg._final_arrow_types(
|
|
output_name,
|
|
input_schema=input_schema,
|
|
)
|
|
)
|
|
return types
|
|
|
|
def _final_cudf_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
"""Return the cuDF dtypes for the final output columns of the GPU aggregation plan.
|
|
|
|
This provides a fallback schema derived from the supplied input schema and runtime input
|
|
schema,
|
|
"""
|
|
input_schema = self._effective_input_schema(input_schema)
|
|
dtypes: Dict[str, Optional[DataType]] = {}
|
|
|
|
if not self._is_global:
|
|
for column in self._shuffle_key_columns:
|
|
dtype = self._effective_column_dtype(column, input_schema)
|
|
if dtype is None:
|
|
dtype = _cudf_column_dtype(df, column)
|
|
elif dtype.is_null_type():
|
|
dtype = DataType.from_numpy("float64")
|
|
dtypes[column] = dtype
|
|
|
|
for agg, output_name, accumulator_prefix in zip(
|
|
self._gpu_aggregates, self._output_names, self._accumulator_prefixes
|
|
):
|
|
dtypes.update(
|
|
{
|
|
column: DataType.from_dtype(dtype)
|
|
for column, dtype in agg._final_cudf_dtypes(
|
|
df,
|
|
output_name,
|
|
accumulator_prefix,
|
|
input_schema=input_schema,
|
|
).items()
|
|
}
|
|
)
|
|
return dtypes
|
|
|
|
def _partial_accumulator_dtypes(
|
|
self,
|
|
df: cudf.DataFrame,
|
|
key_columns: Tuple[str, ...],
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Dict[str, Optional[DataType]]:
|
|
dtypes: Dict[str, Optional[DataType]] = {}
|
|
for column in key_columns:
|
|
if column not in df.columns:
|
|
continue
|
|
dtype = self._effective_column_dtype(column, input_schema)
|
|
if dtype is None:
|
|
dtype = _cudf_column_dtype(df, column)
|
|
elif dtype.is_null_type():
|
|
dtype = DataType.from_numpy("float64")
|
|
dtypes[column] = dtype
|
|
for agg, accumulator_prefix in zip(
|
|
self._gpu_aggregates, self._accumulator_prefixes
|
|
):
|
|
dtypes.update(
|
|
{
|
|
column: DataType.from_dtype(dtype) if dtype is not None else None
|
|
for column, dtype in agg._partial_accumulator_dtypes(
|
|
df,
|
|
accumulator_prefix,
|
|
input_schema=input_schema,
|
|
).items()
|
|
}
|
|
)
|
|
return dtypes
|
|
|
|
|
|
def build_gpu_aggregation_plan(
|
|
key_columns: Tuple[str, ...],
|
|
aggregation_fns: Tuple[Union[AggregateFn, GPUAggregateFn], ...],
|
|
input_schema: Optional[Schema] = None,
|
|
) -> Union[GPUAggregationPlan, str]:
|
|
"""Build a GPU aggregation plan.
|
|
|
|
Args:
|
|
key_columns: The key columns to group by.
|
|
aggregation_fns: The aggregation functions to apply.
|
|
input_schema: The schema of the input data.
|
|
|
|
Returns:
|
|
A GPU aggregation plan if supported, otherwise a fallback reason string.
|
|
"""
|
|
if not aggregation_fns:
|
|
# No aggregation functions, no plan needed.
|
|
return "no aggregation functions were provided."
|
|
|
|
has_gpu_aggregate = any(isinstance(agg, GPUAggregateFn) for agg in aggregation_fns)
|
|
missing_key_columns = [
|
|
column
|
|
for column in key_columns
|
|
if _schema_column_dtype(input_schema, column) is None
|
|
]
|
|
if missing_key_columns and not has_gpu_aggregate:
|
|
# Missing key columns in the input schema, fallback to CPU.
|
|
return (
|
|
"missing input schema for key column(s): "
|
|
f"{', '.join(missing_key_columns)}."
|
|
)
|
|
|
|
gpu_aggregates: List[GPUAggregateFn] = []
|
|
|
|
for agg in aggregation_fns:
|
|
if isinstance(agg, GPUAggregateFn):
|
|
# handle subclasses of GPUAggregateFn as-is (e.g. custom GPU aggregations)
|
|
gpu_aggregate = agg
|
|
else:
|
|
# try to convert built-in GPU aggregation functions to GPU equivalents
|
|
if not isinstance(agg, AggregateFnV2):
|
|
return (
|
|
f"{type(agg).__name__} is not supported by GPU aggregation "
|
|
"because it is not an AggregateFnV2."
|
|
)
|
|
|
|
target_column = agg.get_target_column()
|
|
source_dtype = _schema_column_dtype(input_schema, target_column)
|
|
|
|
if isinstance(agg, Count):
|
|
gpu_aggregate = GPUCount(agg, source_dtype=source_dtype)
|
|
elif target_column is None:
|
|
return (
|
|
f"{type(agg).__name__} is not supported by GPU aggregation "
|
|
"without a target column."
|
|
)
|
|
elif isinstance(agg, Sum):
|
|
gpu_aggregate = GPUSum(agg, source_dtype=source_dtype)
|
|
elif isinstance(agg, Min):
|
|
gpu_aggregate = GPUMin(agg, source_dtype=source_dtype)
|
|
elif isinstance(agg, Max):
|
|
gpu_aggregate = GPUMax(agg, source_dtype=source_dtype)
|
|
elif isinstance(agg, Mean):
|
|
gpu_aggregate = GPUMean(agg, source_dtype=source_dtype)
|
|
else:
|
|
# Any unsupported built-in aggregation in the list falls back
|
|
# the entire list to CPU.
|
|
return f"{type(agg).__name__} is not supported by GPU aggregation."
|
|
|
|
gpu_aggregates.append(gpu_aggregate)
|
|
|
|
return GPUAggregationPlan(
|
|
key_columns,
|
|
tuple(gpu_aggregates),
|
|
accumulator_prefix="__ray_gpu_agg",
|
|
input_schema=input_schema,
|
|
)
|
|
|
|
|
|
@ray.remote(num_gpus=1)
|
|
class GPUHashAggregateActor:
|
|
"""One GPU rank for hash shuffle plus aggregate."""
|
|
|
|
def __init__(
|
|
self,
|
|
nranks: int,
|
|
total_nparts: int,
|
|
aggregation_plan: GPUAggregationPlan,
|
|
rmm_pool_size: Optional[int | str] = None,
|
|
spill_memory_limit: Optional[int | str] = "auto",
|
|
) -> None:
|
|
from ray.data._internal.gpu_shuffle.rapidsmpf_backend import (
|
|
BulkRapidsMPFShuffler,
|
|
)
|
|
|
|
self._aggregation_plan = aggregation_plan
|
|
self._shuffler = BulkRapidsMPFShuffler(
|
|
nranks=nranks,
|
|
total_nparts=total_nparts,
|
|
shuffle_on=list(aggregation_plan.shuffle_key_columns),
|
|
rmm_pool_size=rmm_pool_size,
|
|
spill_memory_limit=spill_memory_limit,
|
|
)
|
|
self._shuffle_columns: Optional[List[str]] = None
|
|
self._runtime_input_schema: Optional[pa.Schema] = (
|
|
aggregation_plan._input_schema
|
|
if isinstance(aggregation_plan._input_schema, pa.Schema)
|
|
else None
|
|
)
|
|
|
|
def setup_root(self) -> Tuple[int, bytes]:
|
|
logger.info("UCXX setup_root starting on GPU hash aggregate rank 0.")
|
|
t0 = time.perf_counter()
|
|
result = self._shuffler.setup_root()
|
|
elapsed = time.perf_counter() - t0
|
|
logger.info(
|
|
"UCXX setup_root completed in %.2fs for GPU hash aggregate rank %d.",
|
|
elapsed,
|
|
result[0],
|
|
)
|
|
return result
|
|
|
|
def setup_worker(self, root_address: bytes) -> None:
|
|
logger.info(
|
|
"UCXX setup_worker starting for GPU hash aggregate "
|
|
"(root_address=%d bytes).",
|
|
len(root_address),
|
|
)
|
|
t0 = time.perf_counter()
|
|
self._shuffler.setup_worker(root_address)
|
|
elapsed = time.perf_counter() - t0
|
|
logger.info("UCXX setup_worker completed in %.2fs.", elapsed)
|
|
|
|
def insert_batch(self, block: Block) -> int:
|
|
import cudf
|
|
|
|
table = BlockAccessor.for_block(block).to_arrow()
|
|
required_columns = self._aggregation_plan.required_columns
|
|
if required_columns:
|
|
projected_table = table.select(list(required_columns))
|
|
df = cudf.DataFrame.from_arrow(projected_table)
|
|
else:
|
|
df = cudf.DataFrame(index=range(table.num_rows))
|
|
|
|
self._runtime_input_schema = self._aggregation_plan.merge_input_schema(
|
|
self._runtime_input_schema,
|
|
table.schema,
|
|
)
|
|
partial = self._aggregation_plan.partial_aggregate(
|
|
df,
|
|
input_schema=self._runtime_input_schema,
|
|
)
|
|
if self._shuffle_columns is None:
|
|
self._shuffle_columns = list(partial.columns)
|
|
|
|
self._shuffler.insert_chunk(table=partial, column_names=self._shuffle_columns)
|
|
return table.num_rows
|
|
|
|
def finish_and_extract(self) -> Iterator[pa.Table | bytes]:
|
|
self._shuffler.insert_finished()
|
|
|
|
import cudf
|
|
from rapidsmpf.utils.cudf import pylibcudf_to_cudf_dataframe
|
|
|
|
self._shuffle_columns = self._shuffle_columns or list(
|
|
self._aggregation_plan.shuffle_key_columns
|
|
+ self._aggregation_plan.accumulator_columns
|
|
)
|
|
|
|
for partition_id, partition in self._shuffler.extract():
|
|
exec_stats_builder = BlockExecStats.builder()
|
|
if partition.num_columns() == 0:
|
|
cdf = cudf.DataFrame()
|
|
else:
|
|
cdf = pylibcudf_to_cudf_dataframe(
|
|
partition, column_names=self._shuffle_columns
|
|
).copy(deep=True)
|
|
|
|
output_df = self._aggregation_plan.final_aggregate(
|
|
cdf,
|
|
input_schema=self._runtime_input_schema,
|
|
)
|
|
block = output_df.to_arrow(preserve_index=False)
|
|
block = self._aggregation_plan.normalize_output_arrow(
|
|
block, input_schema=self._runtime_input_schema
|
|
)
|
|
|
|
existing_metadata = block.schema.metadata or {}
|
|
tagged_schema = block.schema.with_metadata(
|
|
{**existing_metadata, _GPU_PARTITION_ID_KEY: str(partition_id).encode()}
|
|
)
|
|
exec_stats = exec_stats_builder.build()
|
|
stats = yield block
|
|
if stats:
|
|
object.__setattr__(
|
|
exec_stats, "block_ser_time_s", stats.object_creation_dur_s
|
|
)
|
|
block_meta = BlockMetadataWithSchema.from_block(
|
|
block, block_exec_stats=exec_stats
|
|
)
|
|
bm = BlockMetadataWithSchema.from_metadata(
|
|
block_meta.metadata, schema=tagged_schema
|
|
)
|
|
yield pickle.dumps(bm)
|
|
|
|
|
|
class GPUHashAggregateOperator(GPUShuffleOperator):
|
|
"""GPU-native hash aggregate using RAPIDS MPF for the shuffle stage."""
|
|
|
|
def __init__(
|
|
self,
|
|
data_context: DataContext,
|
|
input_op: PhysicalOperator,
|
|
key_columns: Tuple[str, ...],
|
|
aggregation_plan: GPUAggregationPlan,
|
|
*,
|
|
num_partitions: Optional[int] = None,
|
|
) -> None:
|
|
if aggregation_plan is None:
|
|
raise ValueError(
|
|
"GPUHashAggregateOperator received unsupported aggregations."
|
|
)
|
|
|
|
nranks = _derive_num_gpu_ranks(data_context)
|
|
if len(key_columns) == 0:
|
|
# global aggregation
|
|
target_num_partitions = 1
|
|
elif num_partitions is not None:
|
|
# user-specified number of partitions
|
|
target_num_partitions = num_partitions
|
|
else:
|
|
# estimate number of partitions from input operator, otherwise use default
|
|
input_logical_op = input_op._logical_operators[0]
|
|
target_num_partitions = (
|
|
input_logical_op.estimated_num_outputs()
|
|
or data_context.default_hash_shuffle_parallelism
|
|
)
|
|
# rapidsmpf requires total_nparts >= nranks
|
|
target_num_partitions = max(target_num_partitions, nranks)
|
|
|
|
rank_pool = GPURankPool(
|
|
nranks=nranks,
|
|
total_nparts=target_num_partitions,
|
|
setup_timeout_s=data_context.gpu_shuffle_setup_timeout_s,
|
|
actor_cls_factory=lambda: GPUHashAggregateActor,
|
|
actor_kwargs={
|
|
"aggregation_plan": aggregation_plan,
|
|
"rmm_pool_size": data_context.gpu_shuffle_rmm_pool_size,
|
|
"spill_memory_limit": data_context.gpu_shuffle_spill_memory_limit,
|
|
},
|
|
log_label="GPUHashAggregatePool",
|
|
label_selector=data_context.execution_options.label_selector,
|
|
)
|
|
|
|
super().__init__(
|
|
input_op,
|
|
data_context,
|
|
key_columns=aggregation_plan.shuffle_key_columns,
|
|
columns=None,
|
|
num_partitions=target_num_partitions,
|
|
should_sort=False,
|
|
name=(
|
|
f"GPUHashAggregate(key_columns={key_columns}, "
|
|
f"num_partitions={target_num_partitions})"
|
|
),
|
|
nranks=nranks,
|
|
rank_pool=rank_pool,
|
|
)
|
|
|
|
self._aggregation_plan = aggregation_plan
|
|
|
|
def get_sub_progress_bar_names(self) -> List[str]:
|
|
return ["GPU Shuffle", "GPU Aggregation"]
|
|
|
|
def set_sub_progress_bar(self, name: str, pg: "BaseProgressBar") -> None:
|
|
if name == "GPU Shuffle":
|
|
self._shuffle_bar = pg
|
|
elif name == "GPU Aggregation":
|
|
self._reduce_bar = pg
|
|
|
|
def get_stats(self) -> Dict[str, List[BlockStats]]:
|
|
shuffle_name = f"{self._name}_shuffle"
|
|
aggregate_name = f"{self._name}_aggregate"
|
|
return {
|
|
shuffle_name: self._shuffled_blocks_stats,
|
|
aggregate_name: self._output_blocks_stats,
|
|
}
|