wehub-resource-sync
265 lines
8.0 KiB
Python
265 lines
8.0 KiB
Python
"""Utility functions for expression-based operations."""
|
|
|
|
from typing import TYPE_CHECKING, Any, Callable, Hashable, List, Optional
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data.expressions import (
|
|
AliasExpr,
|
|
BinaryExpr,
|
|
ColumnExpr,
|
|
DownloadExpr,
|
|
Expr,
|
|
LiteralExpr,
|
|
MonotonicallyIncreasingIdExpr,
|
|
RandomExpr,
|
|
UDFExpr,
|
|
UnaryExpr,
|
|
UUIDExpr,
|
|
)
|
|
|
|
|
|
def _get_setting_with_copy_warning() -> Optional[type]:
|
|
"""Get the SettingWithCopyWarning class from pandas, if available.
|
|
|
|
Pandas has moved/renamed this warning across versions, and pandas 3.x may not
|
|
expose it at all. This function handles the version differences gracefully
|
|
using hasattr checks instead of try-except blocks.
|
|
|
|
Returns:
|
|
The SettingWithCopyWarning class if found, None otherwise.
|
|
"""
|
|
import pandas as pd
|
|
|
|
# Use hasattr to avoid try-catch blocks as suggested
|
|
if hasattr(pd.core.common, "SettingWithCopyWarning"):
|
|
return pd.core.common.SettingWithCopyWarning
|
|
elif hasattr(pd.errors, "SettingWithCopyWarning"):
|
|
return pd.errors.SettingWithCopyWarning
|
|
else:
|
|
# Warning not available in this pandas version
|
|
return None
|
|
|
|
|
|
def _create_callable_class_udf_init_fn(
|
|
exprs: List["Expr"],
|
|
) -> Optional[Callable[[], None]]:
|
|
"""Create an init_fn to initialize all callable class UDFs in expressions.
|
|
|
|
This function collects all _CallableClassUDF instances from the given expressions,
|
|
groups them by their callable_class_spec key, and returns an init_fn that
|
|
initializes each group at actor startup. UDFs with the same key (same class and
|
|
constructor args) share a single instance to ensure all are properly initialized.
|
|
|
|
Args:
|
|
exprs: List of expressions to collect callable class UDFs from.
|
|
|
|
Returns:
|
|
An init_fn that initializes all callable class UDFs, or None if there are
|
|
no callable class UDFs in the expressions.
|
|
"""
|
|
from ray.data._internal.planner.plan_expression.expression_visitors import (
|
|
_CallableClassUDFCollector,
|
|
)
|
|
|
|
callable_class_udfs = []
|
|
for expr in exprs:
|
|
collector = _CallableClassUDFCollector()
|
|
collector.visit(expr)
|
|
callable_class_udfs.extend(collector.get_callable_class_udfs())
|
|
|
|
if not callable_class_udfs:
|
|
return None
|
|
|
|
# Group UDFs by callable_class_spec key.
|
|
# Multiple _CallableClassUDF objects may have the same key (same class + args).
|
|
# We need to initialize ALL of them, sharing a single instance per key.
|
|
udfs_by_key = {}
|
|
for udf in callable_class_udfs:
|
|
key = udf.callable_class_spec.make_key()
|
|
if key not in udfs_by_key:
|
|
udfs_by_key[key] = []
|
|
udfs_by_key[key].append(udf)
|
|
|
|
def init_fn():
|
|
for udfs_with_same_key in udfs_by_key.values():
|
|
# Initialize the first UDF to create the instance
|
|
first_udf = udfs_with_same_key[0]
|
|
first_udf.init()
|
|
# Share the instance with all other UDFs that have the same key
|
|
for other_udf in udfs_with_same_key[1:]:
|
|
other_udf._instance = first_udf._instance
|
|
|
|
return init_fn
|
|
|
|
|
|
def _call_udf_instance_with_async_bridge(
|
|
instance: Any,
|
|
*args,
|
|
**kwargs,
|
|
) -> Any:
|
|
"""Call a UDF instance, bridging from sync context to async if needed.
|
|
|
|
This handles the complexity of calling callable class UDF instances that may
|
|
be sync, async coroutine, or async generator functions.
|
|
|
|
Args:
|
|
instance: The callable instance to call
|
|
*args: Positional arguments
|
|
**kwargs: Keyword arguments
|
|
|
|
Returns:
|
|
The result of calling the instance
|
|
"""
|
|
import asyncio
|
|
import inspect
|
|
|
|
# Check if the instance's __call__ is async
|
|
if inspect.iscoroutinefunction(instance.__call__):
|
|
# Async coroutine: bridge from sync to async
|
|
return asyncio.run(instance(*args, **kwargs))
|
|
elif inspect.isasyncgenfunction(instance.__call__):
|
|
# Async generator: collect results
|
|
async def _collect():
|
|
results = []
|
|
async for item in instance(*args, **kwargs):
|
|
results.append(item)
|
|
# In expressions, the UDF must return a single array with the same
|
|
# length as the input (one output element per input row).
|
|
# If the async generator yields multiple arrays, we take the last one
|
|
# since expressions don't support multi-batch output semantics.
|
|
if not results:
|
|
return None
|
|
elif len(results) == 1:
|
|
return results[0]
|
|
else:
|
|
import logging
|
|
|
|
logging.warning(
|
|
f"Async generator yielded {len(results)} values in expression context; "
|
|
"only the last (most recent) is returned. Use map_batches for multi-yield support."
|
|
)
|
|
return results[-1]
|
|
|
|
return asyncio.run(_collect())
|
|
else:
|
|
# Synchronous instance - direct call
|
|
return instance(*args, **kwargs)
|
|
|
|
|
|
def _make_hashable(value: Any) -> Hashable:
|
|
try:
|
|
hash(value)
|
|
return value
|
|
except TypeError:
|
|
pass
|
|
|
|
if isinstance(value, list):
|
|
return tuple(_make_hashable(v) for v in value)
|
|
if isinstance(value, tuple):
|
|
return tuple(_make_hashable(v) for v in value)
|
|
if isinstance(value, dict):
|
|
return tuple(
|
|
sorted(
|
|
((k, _make_hashable(v)) for k, v in value.items()),
|
|
key=lambda item: repr(item[0]),
|
|
)
|
|
)
|
|
if isinstance(value, set):
|
|
return frozenset(_make_hashable(v) for v in value)
|
|
|
|
return repr(value)
|
|
|
|
|
|
def _data_type_key(expr: "Expr") -> Hashable:
|
|
return repr(getattr(expr, "data_type", None))
|
|
|
|
|
|
def _udf_function_key(fn: Any) -> Hashable:
|
|
from ray.data.expressions import _CallableClassUDF
|
|
|
|
if isinstance(fn, _CallableClassUDF):
|
|
return ("callable_class", fn.callable_class_spec.make_key())
|
|
return ("function", _make_hashable(fn))
|
|
|
|
|
|
def _column_fingerprint_key(expr: "ColumnExpr") -> Hashable:
|
|
return ("column", expr.name)
|
|
|
|
|
|
def _literal_fingerprint_key(expr: "LiteralExpr") -> Hashable:
|
|
return ("literal", type(expr.value), _make_hashable(expr.value))
|
|
|
|
|
|
def _binary_fingerprint_key(
|
|
expr: "BinaryExpr", left_key: Hashable, right_key: Hashable
|
|
) -> Hashable:
|
|
return ("binary", expr.op, left_key, right_key)
|
|
|
|
|
|
def _unary_fingerprint_key(expr: "UnaryExpr", operand_key: Hashable) -> Hashable:
|
|
return ("unary", expr.op, operand_key)
|
|
|
|
|
|
def _udf_fingerprint_key(
|
|
expr: "UDFExpr",
|
|
arg_keys: tuple[Hashable, ...],
|
|
kwarg_keys: tuple[tuple[str, Hashable], ...],
|
|
) -> Hashable:
|
|
from ray.data.expressions import PyArrowComputeUDFExpr
|
|
|
|
if isinstance(expr, PyArrowComputeUDFExpr):
|
|
return (
|
|
"pyarrow_compute_udf",
|
|
_make_hashable(expr.pc_func),
|
|
_make_hashable(expr.pc_positional),
|
|
_make_hashable(expr.pc_kwargs),
|
|
arg_keys,
|
|
kwarg_keys,
|
|
_data_type_key(expr),
|
|
)
|
|
|
|
return (
|
|
"udf",
|
|
_udf_function_key(expr.fn),
|
|
arg_keys,
|
|
kwarg_keys,
|
|
_data_type_key(expr),
|
|
)
|
|
|
|
|
|
def _alias_fingerprint_key(expr: "AliasExpr", child_key: Hashable) -> Hashable:
|
|
return (
|
|
"alias",
|
|
expr.name,
|
|
expr._is_rename,
|
|
child_key,
|
|
_data_type_key(expr),
|
|
)
|
|
|
|
|
|
def _download_fingerprint_key(expr: "DownloadExpr") -> Hashable:
|
|
return ("download", expr.uri_column_name)
|
|
|
|
|
|
def _star_fingerprint_key() -> Hashable:
|
|
return ("star",)
|
|
|
|
|
|
def _monotonically_increasing_id_fingerprint_key(
|
|
expr: "MonotonicallyIncreasingIdExpr",
|
|
) -> Hashable:
|
|
return ("monotonically_increasing_id", expr._instance_id)
|
|
|
|
|
|
def _random_fingerprint_key(expr: "RandomExpr") -> Hashable:
|
|
return (
|
|
"random",
|
|
expr.seed,
|
|
expr.reseed_after_execution,
|
|
_data_type_key(expr),
|
|
)
|
|
|
|
|
|
def _uuid_fingerprint_key(expr: "UUIDExpr") -> Hashable:
|
|
return ("uuid", _data_type_key(expr))
|