| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263 |
- 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
|
