63 lines
2 KiB
Python
63 lines
2 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from pandas import (
|
|
DataFrame,
|
|
Series,
|
|
)
|
|
from pandas.api.executors import BaseExecutionEngine
|
|
|
|
|
|
class MockExecutionEngine(BaseExecutionEngine):
|
|
"""
|
|
Execution Engine to test if the execution engine interface receives and
|
|
uses all parameters provided by the user.
|
|
|
|
Making this engine work as the default Python engine by calling it, no extra
|
|
functionality is implemented here.
|
|
|
|
When testing, this will be called when this engine is provided, and then the
|
|
same pandas.map and pandas.apply function will be called, but without engine,
|
|
executing the default behavior from the python engine.
|
|
"""
|
|
|
|
def map(data, func, args, kwargs, decorator, skip_na):
|
|
kwargs_to_pass = kwargs if isinstance(data, DataFrame) else {}
|
|
return data.map(func, na_action="ignore" if skip_na else None, **kwargs_to_pass)
|
|
|
|
def apply(data, func, args, kwargs, decorator, axis):
|
|
if isinstance(data, Series):
|
|
return data.apply(func, convert_dtype=True, args=args, by_row=False)
|
|
elif isinstance(data, DataFrame):
|
|
return data.apply(
|
|
func,
|
|
axis=axis,
|
|
raw=False,
|
|
result_type=None,
|
|
args=args,
|
|
by_row="compat",
|
|
**kwargs,
|
|
)
|
|
else:
|
|
assert isinstance(data, np.ndarray)
|
|
|
|
def wrap_function(func):
|
|
# https://github.com/numpy/numpy/issues/8352
|
|
def wrapper(*args, **kwargs):
|
|
result = func(*args, **kwargs)
|
|
if isinstance(result, str):
|
|
result = np.array(result, dtype=object)
|
|
return result
|
|
|
|
return wrapper
|
|
|
|
return np.apply_along_axis(wrap_function(func), axis, data, *args, **kwargs)
|
|
|
|
|
|
class MockEngineDecorator:
|
|
__pandas_udf__ = MockExecutionEngine
|
|
|
|
|
|
@pytest.fixture(params=[None, MockEngineDecorator])
|
|
def engine(request):
|
|
return request.param
|