Repository: modin-project/modin
Pull request #7488 (2025-04-01 21:00:54)

FEAT-#7475: Choose the correct __init__ method from extensions and apply casting to __init__.


- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [ ] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [ ] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [ ] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [ ] Resolves #? <!-- issue must be created for each patch -->
- [ ] tests added and passing
- [ ] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->



Patch:
diff --git a/modin/core/storage_formats/pandas/query_compiler_caster.py b/modin/core/storage_formats/pandas/query_compiler_caster.py
index 9f5ef864..7ef0a96e 100644
--- a/modin/core/storage_formats/pandas/query_compiler_caster.py
+++ b/modin/core/storage_formats/pandas/query_compiler_caster.py
@@ -29,6 +29,7 @@ from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union, ValuesV
 from pandas.core.indexes.frozen import FrozenList
 from typing_extensions import Self
 
+from modin.config import Backend
 from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler
 from modin.core.storage_formats.base.query_compiler_calculator import (
     BackendCostCalculator,
@@ -53,11 +54,6 @@ _NON_EXTENDABLE_ATTRIBUTES = {
     "_query_compiler",
     "_get_query_compiler",
     "_copy_into",
-    # TODO(https://github.com/modin-project/modin/issues/7475): Choose the
-    # correct __init__ implementation from extensions. We have to handle
-    # __init__ differently because the object passed to __init__ does not yet
-    # have a query compiler.
-    "__init__",
 }
 
 
@@ -85,14 +81,15 @@ class QueryCompilerCaster(ABC):
         apply_argument_cast_to_class(cls)
 
     @abstractmethod
-    def _get_query_compiler(self) -> BaseQueryCompiler:
+    def _get_query_compiler(self) -> Optional[BaseQueryCompiler]:
         """
         Get the query compiler storing data for this object.
 
         Returns
         -------
-        BaseQueryCompiler
-            The query compiler storing data for this object.
+        Optional[BaseQueryCompiler]
+            The query compiler storing data for this object, if it exists.
+            Otherwise, None.
         """
         pass
 
@@ -337,21 +334,29 @@ def wrap_function_in_argument_caster(
         calculator: BackendCostCalculator = BackendCostCalculator()
 
         def register_query_compilers(arg):
-            if isinstance(arg, QueryCompilerCaster):
-                calculator.add_query_compiler(arg._get_query_compiler())
+            if (
+                isinstance(arg, QueryCompilerCaster)
+                and (qc := arg._get_query_compiler()) is not None
+            ):
+                calculator.add_query_compiler(qc)
+            elif isinstance(arg, BaseQueryCompiler):
+                # We might get query compiler arguments in __init__()
+                calculator.add_query_compiler(arg)
             return arg
 
         visit_nested_args(args, register_query_compilers)
         visit_nested_args(kwargs, register_query_compilers)
 
         # Sometimes the function takes no query compilers as inputs,
-        # e.g. pd.to_datetime(0)
+        # e.g. pd.to_datetime(0), or some cases of DataFrame.__init__()
         if len(calculator._qc_list) > 0:
             result_backend = calculator.calculate()
 
             def cast_to_qc(arg):
                 if not (
                     isinstance(arg, QueryCompilerCaster)
+                    and (original_query_compiler := arg._get_query_compiler())
+                    is not None
                     and arg.get_backend() != result_backend
                 ):
                     return arg
@@ -359,7 +364,7 @@ def wrap_function_in_argument_caster(
                 inplace_update_trackers.append(
                     InplaceUpdateTracker(
                         input_castable=arg,
-                        original_query_compiler=arg._get_query_compiler(),
+                        original_query_compiler=original_query_compiler,
                         new_castable=cast,
                     )
                 )
@@ -367,6 +372,8 @@ def wrap_function_in_argument_caster(
 
             args = visit_nested_args(args, cast_to_qc)
             kwargs = visit_nested_args(kwargs, cast_to_qc)
+        else:
+            result_backend = Backend.get()
 
         if wrapping_function_type is None:
             # TODO(https://github.com/modin-project/modin/issues/7474): dispatch
@@ -381,9 +388,8 @@ def wrap_function_in_argument_caster(
         # When python invokes a method on an object, it passes the object as
         # the first positional argument.
         self = args[0]
-        self_backend = self.get_backend()
-        if name in cls._extensions[self_backend]:
-            f_to_apply = cls._extensions[self_backend][name]
+        if name in cls._extensions[result_backend]:
+            f_to_apply = cls._extensions[result_backend][name]
         else:
             if name not in cls._extensions[None]:
                 raise AttributeError(
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
index 37a07b60..888db2a4 100644
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -4631,6 +4631,6 @@ class BasePandasDataset(QueryCompilerCaster, ClassLogger):
         return sentinel
 
     @disable_logging
-    @_inherit_docstrings(QueryCompilerCaster._copy_into)
+    @_inherit_docstrings(QueryCompilerCaster._get_query_compiler)
     def _get_query_compiler(self):
-        return self._query_compiler
+        return getattr(self, "_query_compiler", None)
diff --git a/modin/pandas/indexing.py b/modin/pandas/indexing.py
index 10913515..67df6e55 100644
--- a/modin/pandas/indexing.py
+++ b/modin/pandas/indexing.py
@@ -298,7 +298,7 @@ class _LocationIndexerBase(QueryCompilerCaster, ClassLogger):
     @disable_logging
     @_inherit_docstrings(QueryCompilerCaster._get_query_compiler)
     def _get_query_compiler(self):
-        return self.qc
+        return getattr(self, "qc", None)
 
     @disable_logging
     @_inherit_docstrings(QueryCompilerCaster.get_backend)


Test patch:
diff --git a/modin/tests/pandas/extensions/test_dataframe_extensions.py b/modin/tests/pandas/extensions/test_dataframe_extensions.py
index 7d811b2f..2cb09ee7 100644
--- a/modin/tests/pandas/extensions/test_dataframe_extensions.py
+++ b/modin/tests/pandas/extensions/test_dataframe_extensions.py
@@ -17,8 +17,12 @@ from unittest import mock
 import pytest
 
 import modin.pandas as pd
+from modin.config import Backend
+from modin.config import context as config_context
 from modin.pandas.api.extensions import register_dataframe_accessor
 
+default___init__ = pd.DataFrame._extensions[None]["__init__"]
+
 
 def test_dataframe_extension_simple_method(Backend1):
     expected_string_val = "Some string value"
@@ -76,6 +80,55 @@ def test_dataframe_extension_method_uses_superclass_method(Backend1):
     assert df.set_backend(Backend1).sort_values(by=0).iloc[0, 0] == 3
 
 
+class TestOverride__init__:
+    def test_override_one_backend_and_pass_no_query_compilers(self):
+        default_backend = Backend.get()
+        backend_init = mock.Mock(wraps=default___init__)
+        register_dataframe_accessor(name="__init__", backend=default_backend)(
+            backend_init
+        )
+        output_df = pd.DataFrame([1], index=["a"], columns=["b"])
+        assert output_df.get_backend() == default_backend
+        backend_init.assert_has_calls(
+            [
+                mock.call(output_df, [1], index=["a"], columns=["b"]),
+                # There's a second, internal call to the dataframe constructor that
+                # uses a different dataframe as `self`.
+                mock.call(mock.ANY, query_compiler=output_df._query_compiler),
+            ]
+        )
+
+    def test_override_one_backend_and_pass_query_compiler_kwarg(self):
+        backend = "Pandas"
+        backend_init = mock.Mock(wraps=default___init__)
+        register_dataframe_accessor(name="__init__", backend=backend)(backend_init)
+
+        with config_context(Backend=backend):
+            input_df = pd.DataFrame()
+
+        backend_init.reset_mock()
+        output_df = pd.DataFrame(query_compiler=input_df._query_compiler)
+        assert output_df.get_backend() == backend
+        backend_init.assert_called_once_with(
+            output_df, query_compiler=input_df._query_compiler
+        )
+
+    @pytest.mark.parametrize("input_backend", ["Python_Test", "Pandas"])
+    def test_override_all_backends_and_pass_query_compiler_kwarg(self, input_backend):
+        backend_init = mock.Mock(wraps=default___init__)
+        register_dataframe_accessor(name="__init__")(backend_init)
+
+        with config_context(Backend=input_backend):
+            input_df = pd.DataFrame()
+
+        backend_init.reset_mock()
+        output_df = pd.DataFrame(query_compiler=input_df._query_compiler)
+        assert output_df.get_backend() == input_backend
+        backend_init.assert_called_once_with(
+            output_df, query_compiler=input_df._query_compiler
+        )
+
+
 class TestDunders:
     """
     Make sure to test that we override special "dunder" methods like __len__
diff --git a/modin/tests/pandas/extensions/test_series_extensions.py b/modin/tests/pandas/extensions/test_series_extensions.py
index 5fbea58f..33f717fd 100644
--- a/modin/tests/pandas/extensions/test_series_extensions.py
+++ b/modin/tests/pandas/extensions/test_series_extensions.py
@@ -12,12 +12,17 @@
 # governing permissions and limitations under the License.
 
 import re
+from unittest import mock
 
 import pytest
 
 import modin.pandas as pd
+from modin.config import Backend
+from modin.config import context as config_context
 from modin.pandas.api.extensions import register_series_accessor
 
+default___init__ = pd.Series._extensions[None]["__init__"]
+
 
 def test_series_extension_simple_method(Backend1):
     expected_string_val = "Some string value"
@@ -75,6 +80,49 @@ def test_series_extension_method_uses_superclass_method(Backend1):
     assert series.set_backend(Backend1).sort_values().iloc[0] == 3
 
 
+class TestOverride__init__:
+    def test_override_one_backend_and_pass_no_query_compilers(self):
+        default_backend = Backend.get()
+        backend_init = mock.Mock(wraps=default___init__)
+        register_series_accessor(name="__init__", backend=default_backend)(backend_init)
+        output_series = pd.Series([1], index=["a"])
+        assert output_series.get_backend() == default_backend
+        backend_init.assert_has_calls(
+            [
+                mock.call(output_series, [1], index=["a"]),
+            ]
+        )
+
+    def test_override_one_backend_and_pass_query_compiler_kwarg(self):
+        backend_init = mock.Mock(wraps=default___init__)
+        register_series_accessor(name="__init__", backend="Pandas")(backend_init)
+
+        with config_context(Backend="Pandas"):
+            input_series = pd.Series()
+
+        backend_init.reset_mock()
+        output_series = pd.Series(query_compiler=input_series._query_compiler)
+        assert output_series.get_backend() == "Pandas"
+        backend_init.assert_called_once_with(
+            output_series, query_compiler=input_series._query_compiler
+        )
+
+    @pytest.mark.parametrize("input_backend", ["Python_Test", "Pandas"])
+    def test_override_all_backends_and_pass_query_compiler_kwarg(self, input_backend):
+        backend_init = mock.Mock(wraps=default___init__)
+        register_series_accessor(name="__init__")(backend_init)
+
+        with config_context(Backend=input_backend):
+            input_series = pd.Series()
+
+        backend_init.reset_mock()
+        output_series = pd.Series(query_compiler=input_series._query_compiler)
+        assert output_series.get_backend() == input_backend
+        backend_init.assert_called_once_with(
+            output_series, query_compiler=input_series._query_compiler
+        )
+
+
 class TestDunders:
     """
     Make sure to test that we override special "dunder" methods like __len__
diff --git a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
index 3db1a055..4e6e950f 100644
--- a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
+++ b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
@@ -195,13 +195,20 @@ def test_two_same_backend(pico_df):
     assert df3.get_backend() == "Pico"
 
 
-def test_cast_to_second_backend(pico_df, cluster_df):
+def test_cast_to_second_backend_with_concat(pico_df, cluster_df):
     df3 = pd.concat([pico_df, cluster_df], axis=1)
     assert pico_df.get_backend() == "Pico"
     assert cluster_df.get_backend() == "Cluster"
     assert df3.get_backend() == "Cluster"  # result should be on cluster
 
 
+def test_cast_to_second_backend_with___init__(pico_df, cluster_df):
+    df3 = pd.DataFrame({"pico": pico_df.iloc[:, 0], "cluster": cluster_df.iloc[:, 0]})
+    assert pico_df.get_backend() == "Pico"
+    assert cluster_df.get_backend() == "Cluster"
+    assert df3.get_backend() == "Cluster"  # result should be on cluster
+
+
 def test_cast_to_first_backend(pico_df, cluster_df):
     df3 = pd.concat([cluster_df, pico_df], axis=1)
     assert pico_df.get_backend() == "Pico"
@@ -209,6 +216,18 @@ def test_cast_to_first_backend(pico_df, cluster_df):
     assert df3.get_backend() == cluster_df.get_backend()  # result should be on cluster
 
 
+def test_cast_to_first_backend_with___init__(pico_df, cluster_df):
+    df3 = pd.DataFrame(
+        {
+            "cluster": cluster_df.iloc[:, 0],
+            "pico": pico_df.iloc[:, 0],
+        }
+    )
+    assert pico_df.get_backend() == "Pico"
+    assert cluster_df.get_backend() == "Cluster"
+    assert df3.get_backend() == "Cluster"  # result should be on cluster
+
+
 def test_no_solution(pico_df, local_df, cluster_df, cloud_df):
     with pytest.raises(ValueError, match=r"Pico,Local_machine,Cluster,Cloud"):
         pd.concat(axis=1, objs=[pico_df, local_df, cluster_df, cloud_df])