FEAT-#6767: Implement 'modin.utils.enable_exp_mode' context manager

Signed-off-by: Anatoly Myachev <[email protected]>

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.
 

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(

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.