wehub-resource-sync
319 lines
11 KiB
Python
319 lines
11 KiB
Python
import hashlib
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Callable,
|
|
Deque,
|
|
Dict,
|
|
List,
|
|
Optional,
|
|
Union,
|
|
)
|
|
|
|
import ray
|
|
from ray.air.util.data_batch_conversion import BatchFormat
|
|
from ray.data.aggregate import AggregateFnV2
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data.dataset import Dataset
|
|
|
|
|
|
@DeveloperAPI
|
|
def simple_split_tokenizer(value: str) -> List[str]:
|
|
"""Tokenize a string using a split on spaces."""
|
|
return value.split(" ")
|
|
|
|
|
|
@DeveloperAPI
|
|
def simple_hash(value: object, num_features: int) -> int:
|
|
"""Deterministically hash a value into the integer space."""
|
|
encoded_value = str(value).encode()
|
|
hashed_value = hashlib.sha256(encoded_value)
|
|
hashed_value_int = int(hashed_value.hexdigest(), 16)
|
|
return hashed_value_int % num_features
|
|
|
|
|
|
class BaseStatSpec:
|
|
"""Encapsulates a statistical computation with optional post-processing."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
stat_fn: Union[AggregateFnV2, Callable],
|
|
post_process_fn: Callable = lambda x: x,
|
|
):
|
|
self.stat_fn = stat_fn
|
|
self.post_process_fn = post_process_fn
|
|
|
|
|
|
class AggregateStatSpec(BaseStatSpec):
|
|
"""Represents an AggregateFnV2 spec for a single column."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
aggregator_fn: Union[AggregateFnV2, Callable[[str], AggregateFnV2]],
|
|
post_process_fn: Callable = lambda x: x,
|
|
column: Optional[str] = None,
|
|
batch_format: Optional[BatchFormat] = None,
|
|
):
|
|
super().__init__(
|
|
stat_fn=aggregator_fn,
|
|
post_process_fn=post_process_fn,
|
|
)
|
|
self.column = column
|
|
self.batch_format = batch_format
|
|
|
|
|
|
class CallableStatSpec(BaseStatSpec):
|
|
"""Represents a user-defined stat function that operates outside Dataset.aggregate."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
stat_fn: Callable,
|
|
stat_key_fn: Optional[Callable[[str], str]],
|
|
post_key_fn: Optional[Callable[[str], str]],
|
|
post_process_fn: Callable = lambda x: x,
|
|
columns: List[str],
|
|
):
|
|
super().__init__(
|
|
stat_fn=stat_fn,
|
|
post_process_fn=post_process_fn,
|
|
)
|
|
self.columns = columns
|
|
self.stat_key_fn = stat_key_fn
|
|
self.post_key_fn = post_key_fn
|
|
|
|
|
|
class StatComputationPlan:
|
|
"""
|
|
Encapsulates a set of aggregators (AggregateFnV2) and legacy stat functions
|
|
to compute statistics over a Ray dataset.
|
|
|
|
Supports two types of aggregations:
|
|
1. AggregateFnV2-based aggregators, which are batch-executed using `Dataset.aggregate(...)`.
|
|
2. Callable-based stat functions, executed sequentially (legacy use case).
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._aggregators: Deque[BaseStatSpec] = deque()
|
|
|
|
def reset(self):
|
|
self._aggregators.clear()
|
|
|
|
def add_aggregator(
|
|
self,
|
|
*,
|
|
aggregator_fn: Callable[[str], AggregateFnV2],
|
|
post_process_fn: Callable = lambda x: x,
|
|
columns: List[str],
|
|
batch_format: Optional[BatchFormat] = None,
|
|
) -> None:
|
|
"""
|
|
Registers an AggregateFnV2 factory for one or more columns.
|
|
|
|
Args:
|
|
aggregator_fn: A callable (typically a lambda or class) that accepts a column name and returns an instance of AggregateFnV2.
|
|
The aggregator should set its name using alias_name parameter to control the output key.
|
|
post_process_fn: Function to post-process the aggregated result.
|
|
columns: List of column names to aggregate.
|
|
batch_format: The batch format for aggregation results. If ARROW, results
|
|
are kept in Arrow format for post_process_fn. Otherwise,
|
|
results are converted to Python/pandas format.
|
|
"""
|
|
for column in columns:
|
|
agg_instance = aggregator_fn(column)
|
|
self._aggregators.append(
|
|
AggregateStatSpec(
|
|
aggregator_fn=agg_instance,
|
|
post_process_fn=post_process_fn,
|
|
column=column,
|
|
batch_format=batch_format,
|
|
)
|
|
)
|
|
|
|
def add_callable_stat(
|
|
self,
|
|
*,
|
|
stat_fn: Callable[[], Any],
|
|
stat_key_fn: Callable[[str], str],
|
|
post_key_fn: Optional[Callable[[str], str]] = None,
|
|
post_process_fn: Callable = lambda x: x,
|
|
columns: List[str],
|
|
) -> None:
|
|
"""
|
|
Registers a custom stat function to be run sequentially.
|
|
|
|
This supports legacy use cases where arbitrary callables are needed
|
|
and cannot be run via Dataset.aggregate().
|
|
|
|
Args:
|
|
stat_fn: A zero-argument callable that returns the stat.
|
|
stat_key_fn: A callable that takes a column name and returns the key for the stat.
|
|
post_key_fn: Optional; a callable to post-process the key. If not provided, stat_key_fn is used.
|
|
post_process_fn: Function to post-process the result.
|
|
columns: List of column names to compute the stat for.
|
|
"""
|
|
self._aggregators.append(
|
|
CallableStatSpec(
|
|
stat_fn=stat_fn,
|
|
post_process_fn=post_process_fn,
|
|
columns=columns,
|
|
stat_key_fn=stat_key_fn,
|
|
post_key_fn=post_key_fn or stat_key_fn,
|
|
)
|
|
)
|
|
|
|
def compute(self, dataset: "Dataset") -> Dict[str, Any]:
|
|
"""
|
|
Executes all registered aggregators and stat functions.
|
|
|
|
AggregateFnV2-based aggregators are batched and executed via Dataset.aggregate().
|
|
Callable-based stat functions are run sequentially.
|
|
|
|
Args:
|
|
dataset: The Ray Dataset to compute statistics on.
|
|
|
|
Returns:
|
|
A dictionary of computed statistics.
|
|
"""
|
|
stats = {}
|
|
# Run batched aggregators (AggregateFnV2)
|
|
aggregators = self._get_aggregate_fn_list()
|
|
if aggregators:
|
|
agg_ds = dataset.groupby(None).aggregate(*aggregators)
|
|
arrow_refs = agg_ds.to_arrow_refs()
|
|
if not arrow_refs:
|
|
raise ValueError("Aggregation returned no results")
|
|
arrow_table = ray.get(arrow_refs[0])
|
|
for spec in self._get_aggregate_specs():
|
|
stat_key = spec.stat_fn.name
|
|
# Aggregation returns single row - extract the scalar value
|
|
# ChunkedArray[0] handles multi-chunk arrays automatically
|
|
agg_result = arrow_table.column(stat_key)[0]
|
|
# Convert to appropriate format based on batch_format
|
|
if spec.batch_format == BatchFormat.ARROW:
|
|
# Pass Arrow scalar (e.g., ListScalar) for Arrow-optimized post-processing
|
|
stats[stat_key] = spec.post_process_fn(agg_result)
|
|
else:
|
|
# Convert to Python for pandas-style post-processing
|
|
stats[stat_key] = spec.post_process_fn(agg_result.as_py())
|
|
|
|
# Run sequential stat functions
|
|
for spec in self._get_custom_stat_fn_specs():
|
|
result = spec.stat_fn(spec.stat_key_fn)
|
|
for col in spec.columns:
|
|
stat_key = spec.stat_key_fn(col)
|
|
post_key = spec.post_key_fn(col)
|
|
stats[post_key] = spec.post_process_fn(result[stat_key])
|
|
|
|
return stats
|
|
|
|
def _get_aggregate_fn_list(self) -> List[AggregateFnV2]:
|
|
return [
|
|
spec.stat_fn
|
|
for spec in self._aggregators
|
|
if isinstance(spec, AggregateStatSpec)
|
|
]
|
|
|
|
def _get_aggregate_specs(self) -> List[AggregateStatSpec]:
|
|
return [
|
|
spec for spec in self._aggregators if isinstance(spec, AggregateStatSpec)
|
|
]
|
|
|
|
def _get_custom_stat_fn_specs(self) -> List[CallableStatSpec]:
|
|
return [
|
|
spec for spec in self._aggregators if isinstance(spec, CallableStatSpec)
|
|
]
|
|
|
|
def has_custom_stat_fn(self):
|
|
return len(self._get_custom_stat_fn_specs()) > 0
|
|
|
|
def __iter__(self):
|
|
"""
|
|
Iterates over all AggregatorSpecs.
|
|
"""
|
|
return iter(self._get_aggregate_specs())
|
|
|
|
|
|
def make_post_processor(base_fn, callbacks: List[Callable]):
|
|
"""
|
|
Wraps a base post-processing function with a sequence of callback functions.
|
|
Useful when multiple post-processing steps need to be applied in order.
|
|
"""
|
|
|
|
def wrapper(result):
|
|
processed = base_fn(result)
|
|
for cb in callbacks:
|
|
processed = cb(processed)
|
|
return processed
|
|
|
|
return wrapper
|
|
|
|
|
|
class _Computed:
|
|
"""
|
|
Wraps a factory callable for defaults that must be computed from the object.
|
|
|
|
Plain callable values (e.g. a tokenizer function stored as an attribute)
|
|
must NOT be wrapped — they will be stored as-is.
|
|
"""
|
|
|
|
def __init__(self, factory: Callable[[Any], Any]) -> None:
|
|
self._factory = factory
|
|
|
|
def __call__(self, obj: Any) -> Any:
|
|
return self._factory(obj)
|
|
|
|
|
|
_REQUIRED_FIELD = object() # Sentinel for required fields with no default value
|
|
|
|
|
|
@dataclass
|
|
class _PublicField:
|
|
"""
|
|
Represents a public field that may have been used in older versions of the code.
|
|
If the field's default value is not _REQUIRED_FIELD, it will be used if neither the private field nor the public field is present during unpickling.
|
|
Otherwise, the field is required and must be present as either the private or public field during unpickling, or a ValueError will be raised.
|
|
Used for backwards compatibility during unpickling.
|
|
"""
|
|
|
|
public_field: str
|
|
default: Any = _REQUIRED_FIELD
|
|
|
|
|
|
def migrate_private_fields(
|
|
obj: Any,
|
|
*,
|
|
fields: Dict[str, _PublicField],
|
|
) -> None:
|
|
"""
|
|
Migrates old public field names to new private field names during unpickling for backwards compatibility.
|
|
"""
|
|
for private_field, public_field_obj in fields.items():
|
|
if private_field not in obj.__dict__:
|
|
if public_field_obj.public_field in obj.__dict__:
|
|
# Migrate from old public field names to new private field names
|
|
setattr(
|
|
obj, private_field, obj.__dict__.pop(public_field_obj.public_field)
|
|
)
|
|
elif public_field_obj.default is _REQUIRED_FIELD:
|
|
raise ValueError(
|
|
f"Invalid serialized {type(obj).__name__}: missing required field '{private_field}'."
|
|
)
|
|
else:
|
|
# Set defaults for missing fields.
|
|
# _Computed defaults are called with obj; all other values are stored as-is,
|
|
# including callable objects like tokenizer functions.
|
|
setattr(
|
|
obj,
|
|
private_field,
|
|
public_field_obj.default(obj)
|
|
if isinstance(public_field_obj.default, _Computed)
|
|
else public_field_obj.default,
|
|
)
|