diff --git a/modin/core/execution/dispatching/factories/dispatcher.py b/modin/core/execution/dispatching/factories/dispatcher.py index 623bf354d45..469ed8222ad 100644 --- a/modin/core/execution/dispatching/factories/dispatcher.py +++ b/modin/core/execution/dispatching/factories/dispatcher.py @@ -106,6 +106,12 @@ class FactoryDispatcher(object): __factory: factories.BaseFactory = None + @classmethod + def _reset_and_update_factory(cls): + if cls.__factory is None: + return + cls.__factory = cls._update_factory() + @classmethod def get_factory(cls) -> factories.BaseFactory: """Get current factory.""" @@ -118,8 +124,7 @@ def get_factory(cls) -> factories.BaseFactory: return cls.__factory @classmethod - # FIXME: replace `_` parameter with `*args` - def _update_factory(cls, _): + def _update_factory(cls, *args): """ Update and prepare factory with a new one specified via Modin config. diff --git a/modin/experimental/pandas/__init__.py b/modin/experimental/pandas/__init__.py index 0deba579ecb..1f9b5998c0a 100644 --- a/modin/experimental/pandas/__init__.py +++ b/modin/experimental/pandas/__init__.py @@ -48,6 +48,13 @@ to_pickle_distributed, ) + +class ExpDataFrame(DataFrame): # noqa: F405 + pass + + +DataFrame = ExpDataFrame + setattr(DataFrame, "to_pickle_distributed", to_pickle_distributed) # noqa: F405 warnings.warn( diff --git a/modin/utils.py b/modin/utils.py index c42a0516fa3..8b164a14147 100644 --- a/modin/utils.py +++ b/modin/utils.py @@ -14,6 +14,7 @@ """Collection of general utility functions, mostly for internal use.""" import codecs +import contextlib import functools import importlib import inspect @@ -628,6 +629,20 @@ def execute(*objs: Iterable[Any], trigger_hdk_import: bool = False) -> None: query_compiler.force_import() +@contextlib.contextmanager +def enable_exp_mode(): + from modin.config import IsExperimental + from modin.core.execution.dispatching.factories.dispatcher import FactoryDispatcher + + old_value = IsExperimental.get() + IsExperimental.put(True) + FactoryDispatcher._reset_and_update_factory() + yield + IsExperimental.put(old_value or False) + # cleanup potential side effects from importing `modin.experimental.pandas` + FactoryDispatcher._reset_and_update_factory() + + def wrap_into_list(*args: Any, skipna: bool = True) -> List[Any]: """ Wrap a sequence of passed values in a flattened list.