snowflake-ml-python 1.6.0__py3-none-any.whl → 1.6.2__py3-none-any.whl

snowflake/ml/modeling/framework/base.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 import inspect
 from abc import abstractmethod
-from collections import defaultdict
 from datetime import datetime
 from typing import Any, Dict, Iterable, List, Mapping, Optional, Union, overload
 
@@ -18,6 +17,7 @@ from snowflake.ml._internal.exceptions import (
 )
 from snowflake.ml._internal.lineage import lineage_utils
 from snowflake.ml._internal.utils import identifier, parallelize
+from snowflake.ml.data import data_source
 from snowflake.ml.modeling.framework import _utils
 from snowflake.snowpark import functions as F
 
@@ -246,7 +246,7 @@ class Base:
 
     def get_params(self, deep: bool = True) -> Dict[str, Any]:
         """
-        Get parameters for this transformer.
+        Get the snowflake-ml parameters for this transformer.
 
         Args:
             deep: If True, will return the parameters for this transformer and
@@ -265,13 +265,13 @@ class Base:
                 out[key] = value
         return out
 
-    def set_params(self, **params: Dict[str, Any]) -> None:
+    def set_params(self, **params: Any) -> None:
         """
         Set the parameters of this transformer.
 
-        The method works on simple transformers as well as on nested objects.
-        The latter have parameters of the form ``<component>__<parameter>``
-        so that it's possible to update each component of a nested object.
+        The method works on simple transformers as well as on sklearn compatible pipelines with nested
+        objects, once the transformer has been fit. Nested objects have parameters of the form
+        ``<component>__<parameter>`` so that it's possible to update each component of a nested object.
 
         Args:
             **params: Transformer parameter names mapped to their values.
@@ -283,12 +283,28 @@ class Base:
             # simple optimization to gain speed (inspect is slow)
             return
         valid_params = self.get_params(deep=True)
+        valid_skl_params = {}
+        if hasattr(self, "_sklearn_object") and self._sklearn_object is not None:
+            valid_skl_params = self._sklearn_object.get_params()
 
-        nested_params: Dict[str, Any] = defaultdict(dict)  # grouped by prefix
         for key, value in params.items():
-            key, delim, sub_key = key.partition("__")
-            if key not in valid_params:
-                local_valid_params = self._get_param_names()
+            if valid_params.get("steps"):
+                # Recurse through pipeline steps
+                key, _, sub_key = key.partition("__")
+                for name, nested_object in valid_params["steps"]:
+                    if name == key:
+                        nested_object.set_params(**{sub_key: value})
+
+            elif key in valid_params:
+                setattr(self, key, value)
+                valid_params[key] = value
+            elif key in valid_skl_params:
+                # This dictionary would be empty if the following assert were not true, as specified above.
+                assert hasattr(self, "_sklearn_object") and self._sklearn_object is not None
+                setattr(self._sklearn_object, key, value)
+                valid_skl_params[key] = value
+            else:
+                local_valid_params = self._get_param_names() + list(valid_skl_params.keys())
                 raise exceptions.SnowflakeMLException(
                     error_code=error_codes.INVALID_ARGUMENT,
                     original_exception=ValueError(
@@ -298,15 +314,6 @@ class Base:
                     ),
                 )
 
-            if delim:
-                nested_params[key][sub_key] = value
-            else:
-                setattr(self, key, value)
-                valid_params[key] = value
-
-        for key, sub_params in nested_params.items():
-            valid_params[key].set_params(**sub_params)
-
     def get_sklearn_args(
         self,
         default_sklearn_obj: Optional[object] = None,
@@ -427,6 +434,8 @@ class BaseEstimator(Base):
     def fit(self, dataset: Union[snowpark.DataFrame, pd.DataFrame]) -> "BaseEstimator":
         """Runs universal logics for all fit implementations."""
         data_sources = lineage_utils.get_data_sources(dataset)
+        if not data_sources and isinstance(dataset, snowpark.DataFrame):
+            data_sources = [data_source.DataFrameInfo(dataset.queries["queries"][-1])]
         lineage_utils.set_data_sources(self, data_sources)
         return self._fit(dataset)
 

snowflake/ml/modeling/gaussian_process/gaussian_process_classifier.py

@@ -4,14 +4,12 @@
 #
 import inspect
 import os
-import posixpath
-from typing import Iterable, Optional, Union, List, Any, Dict, Callable, Set
-from typing_extensions import TypeGuard
+from typing import Iterable, Optional, Union, List, Any, Dict, Set, Tuple
 from uuid import uuid4
 
 import cloudpickle as cp
-import pandas as pd
 import numpy as np
+import pandas as pd
 from numpy import typing as npt
 
 
@@ -24,12 +22,11 @@ from snowflake.ml.modeling.framework.base import BaseTransformer, _process_cols
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.exceptions import error_codes, exceptions, modeling_error_messages
 from snowflake.ml._internal.env_utils import SNOWML_SPROC_ENV
-from snowflake.ml._internal.utils import pkg_version_utils, identifier
+from snowflake.ml._internal.utils import identifier
 from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 from snowflake.ml.modeling._internal.model_trainer_builder import ModelTrainerBuilder
 from snowflake.ml.modeling._internal.transformer_protocols import (
-    ModelTransformHandlers,
     BatchInferenceKwargsTypedDict,
     ScoreKwargsTypedDict
 )
@@ -564,12 +561,23 @@ class GaussianProcessClassifier(BaseTransformer):
             autogenerated=self._autogenerated,
             subproject=_SUBPROJECT,
         )
-        output_result, fitted_estimator = model_trainer.train_fit_predict(
-            drop_input_cols=self._drop_input_cols,
-            expected_output_cols_list=(
-                self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
-            ),
+        expected_output_cols = (
+            self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
         )
+        if isinstance(dataset, DataFrame):
+            expected_output_cols, example_output_pd_df = self._align_expected_output(
+                "fit_predict", dataset, expected_output_cols, output_cols_prefix
+            )
+            output_result, fitted_estimator = model_trainer.train_fit_predict(
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_list=expected_output_cols,
+                example_output_pd_df=example_output_pd_df,
+            )
+        else:
+            output_result, fitted_estimator = model_trainer.train_fit_predict(
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_list=expected_output_cols,
+            )
         self._sklearn_object = fitted_estimator
         self._is_fitted = True
         return output_result
@@ -648,12 +656,41 @@ class GaussianProcessClassifier(BaseTransformer):
 
         return rv
 
-    def _align_expected_output_names(
-        self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str
-    ) -> List[str]:
+    def _align_expected_output(
+        self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str,
+    ) -> Tuple[List[str], pd.DataFrame]:
+        """ Run 1 line of data with the desired method, and return one tuple that consists of the output column names 
+        and output dataframe with 1 line.
+        If the method is fit_predict, run 2 lines of data. 
+        """
         # in case the inferred output column names dimension is different 
         # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
-        sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
+
+        # For fit_predict method, a minimum of 2 is required by MinCovDet, BayesianGaussianMixture
+        # so change the minimum of number of rows to 2
+        num_examples = 2
+        statement_params = telemetry.get_function_usage_statement_params(
+            project=_PROJECT,
+            subproject=_SUBPROJECT,
+            function_name=telemetry.get_statement_params_full_func_name(
+                inspect.currentframe(), GaussianProcessClassifier.__class__.__name__
+            ),
+            api_calls=[Session.call],
+            custom_tags={"autogen": True} if self._autogenerated else None,
+        )
+        if output_cols_prefix == "fit_predict_":
+            if hasattr(self._sklearn_object, "n_clusters"):
+                # cluster classes such as BisectingKMeansTest requires # of examples >= n_clusters
+                num_examples = self._sklearn_object.n_clusters
+            elif hasattr(self._sklearn_object, "min_samples"):
+                # OPTICS default min_samples 5, which requires at least 5 lines of data
+                num_examples = self._sklearn_object.min_samples
+            elif hasattr(self._sklearn_object, "n_neighbors") and hasattr(self._sklearn_object, "n_samples"):
+                # LocalOutlierFactor expects n_neighbors <= n_samples
+                num_examples = self._sklearn_object.n_neighbors
+            sample_pd_df = dataset.select(self.input_cols).limit(num_examples).to_pandas(statement_params=statement_params)
+        else:
+            sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas(statement_params=statement_params)
 
         # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
         # seen during the fit.
@@ -665,12 +702,14 @@ class GaussianProcessClassifier(BaseTransformer):
         output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
         if self.sample_weight_col:
             output_df_columns_set -= set(self.sample_weight_col)
+        
         # if the dimension of inferred output column names is correct; use it
         if len(expected_output_cols_list) == len(output_df_columns_set):
-            return expected_output_cols_list
+            return expected_output_cols_list, output_df_pd
         # otherwise, use the sklearn estimator's output
         else:
-            return sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
+            expected_output_cols_list = sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
+            return expected_output_cols_list, output_df_pd[expected_output_cols_list]
 
     @available_if(original_estimator_has_callable("predict_proba"))  # type: ignore[misc]
     @telemetry.send_api_usage_telemetry(
@@ -718,7 +757,7 @@ class GaussianProcessClassifier(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 
@@ -786,7 +825,7 @@ class GaussianProcessClassifier(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
         elif isinstance(dataset, pd.DataFrame):
@@ -849,7 +888,7 @@ class GaussianProcessClassifier(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 
@@ -914,7 +953,7 @@ class GaussianProcessClassifier(BaseTransformer):
                 drop_input_cols = self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 

snowflake/ml/modeling/gaussian_process/gaussian_process_regressor.py

@@ -4,14 +4,12 @@
 #
 import inspect
 import os
-import posixpath
-from typing import Iterable, Optional, Union, List, Any, Dict, Callable, Set
-from typing_extensions import TypeGuard
+from typing import Iterable, Optional, Union, List, Any, Dict, Set, Tuple
 from uuid import uuid4
 
 import cloudpickle as cp
-import pandas as pd
 import numpy as np
+import pandas as pd
 from numpy import typing as npt
 
 
@@ -24,12 +22,11 @@ from snowflake.ml.modeling.framework.base import BaseTransformer, _process_cols
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.exceptions import error_codes, exceptions, modeling_error_messages
 from snowflake.ml._internal.env_utils import SNOWML_SPROC_ENV
-from snowflake.ml._internal.utils import pkg_version_utils, identifier
+from snowflake.ml._internal.utils import identifier
 from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 from snowflake.ml.modeling._internal.model_trainer_builder import ModelTrainerBuilder
 from snowflake.ml.modeling._internal.transformer_protocols import (
-    ModelTransformHandlers,
     BatchInferenceKwargsTypedDict,
     ScoreKwargsTypedDict
 )
@@ -555,12 +552,23 @@ class GaussianProcessRegressor(BaseTransformer):
             autogenerated=self._autogenerated,
             subproject=_SUBPROJECT,
         )
-        output_result, fitted_estimator = model_trainer.train_fit_predict(
-            drop_input_cols=self._drop_input_cols,
-            expected_output_cols_list=(
-                self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
-            ),
+        expected_output_cols = (
+            self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
         )
+        if isinstance(dataset, DataFrame):
+            expected_output_cols, example_output_pd_df = self._align_expected_output(
+                "fit_predict", dataset, expected_output_cols, output_cols_prefix
+            )
+            output_result, fitted_estimator = model_trainer.train_fit_predict(
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_list=expected_output_cols,
+                example_output_pd_df=example_output_pd_df,
+            )
+        else:
+            output_result, fitted_estimator = model_trainer.train_fit_predict(
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_list=expected_output_cols,
+            )
         self._sklearn_object = fitted_estimator
         self._is_fitted = True
         return output_result
@@ -639,12 +647,41 @@ class GaussianProcessRegressor(BaseTransformer):
 
         return rv
 
-    def _align_expected_output_names(
-        self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str
-    ) -> List[str]:
+    def _align_expected_output(
+        self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str,
+    ) -> Tuple[List[str], pd.DataFrame]:
+        """ Run 1 line of data with the desired method, and return one tuple that consists of the output column names 
+        and output dataframe with 1 line.
+        If the method is fit_predict, run 2 lines of data. 
+        """
         # in case the inferred output column names dimension is different 
         # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
-        sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
+
+        # For fit_predict method, a minimum of 2 is required by MinCovDet, BayesianGaussianMixture
+        # so change the minimum of number of rows to 2
+        num_examples = 2
+        statement_params = telemetry.get_function_usage_statement_params(
+            project=_PROJECT,
+            subproject=_SUBPROJECT,
+            function_name=telemetry.get_statement_params_full_func_name(
+                inspect.currentframe(), GaussianProcessRegressor.__class__.__name__
+            ),
+            api_calls=[Session.call],
+            custom_tags={"autogen": True} if self._autogenerated else None,
+        )
+        if output_cols_prefix == "fit_predict_":
+            if hasattr(self._sklearn_object, "n_clusters"):
+                # cluster classes such as BisectingKMeansTest requires # of examples >= n_clusters
+                num_examples = self._sklearn_object.n_clusters
+            elif hasattr(self._sklearn_object, "min_samples"):
+                # OPTICS default min_samples 5, which requires at least 5 lines of data
+                num_examples = self._sklearn_object.min_samples
+            elif hasattr(self._sklearn_object, "n_neighbors") and hasattr(self._sklearn_object, "n_samples"):
+                # LocalOutlierFactor expects n_neighbors <= n_samples
+                num_examples = self._sklearn_object.n_neighbors
+            sample_pd_df = dataset.select(self.input_cols).limit(num_examples).to_pandas(statement_params=statement_params)
+        else:
+            sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas(statement_params=statement_params)
 
         # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
         # seen during the fit.
@@ -656,12 +693,14 @@ class GaussianProcessRegressor(BaseTransformer):
         output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
         if self.sample_weight_col:
             output_df_columns_set -= set(self.sample_weight_col)
+        
         # if the dimension of inferred output column names is correct; use it
         if len(expected_output_cols_list) == len(output_df_columns_set):
-            return expected_output_cols_list
+            return expected_output_cols_list, output_df_pd
         # otherwise, use the sklearn estimator's output
         else:
-            return sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
+            expected_output_cols_list = sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
+            return expected_output_cols_list, output_df_pd[expected_output_cols_list]
 
     @available_if(original_estimator_has_callable("predict_proba"))  # type: ignore[misc]
     @telemetry.send_api_usage_telemetry(
@@ -707,7 +746,7 @@ class GaussianProcessRegressor(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 
@@ -773,7 +812,7 @@ class GaussianProcessRegressor(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
         elif isinstance(dataset, pd.DataFrame):
@@ -836,7 +875,7 @@ class GaussianProcessRegressor(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 
@@ -901,7 +940,7 @@ class GaussianProcessRegressor(BaseTransformer):
                 drop_input_cols = self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 

snowflake/ml/modeling/impute/iterative_imputer.py

@@ -4,14 +4,12 @@
 #
 import inspect
 import os
-import posixpath
-from typing import Iterable, Optional, Union, List, Any, Dict, Callable, Set
-from typing_extensions import TypeGuard
+from typing import Iterable, Optional, Union, List, Any, Dict, Set, Tuple
 from uuid import uuid4
 
 import cloudpickle as cp
-import pandas as pd
 import numpy as np
+import pandas as pd
 from numpy import typing as npt
 
 
@@ -25,12 +23,11 @@ from snowflake.ml.modeling.framework.base import BaseTransformer, _process_cols
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.exceptions import error_codes, exceptions, modeling_error_messages
 from snowflake.ml._internal.env_utils import SNOWML_SPROC_ENV
-from snowflake.ml._internal.utils import pkg_version_utils, identifier
+from snowflake.ml._internal.utils import identifier
 from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 from snowflake.ml.modeling._internal.model_trainer_builder import ModelTrainerBuilder
 from snowflake.ml.modeling._internal.transformer_protocols import (
-    ModelTransformHandlers,
     BatchInferenceKwargsTypedDict,
     ScoreKwargsTypedDict
 )
@@ -599,12 +596,23 @@ class IterativeImputer(BaseTransformer):
             autogenerated=self._autogenerated,
             subproject=_SUBPROJECT,
         )
-        output_result, fitted_estimator = model_trainer.train_fit_predict(
-            drop_input_cols=self._drop_input_cols,
-            expected_output_cols_list=(
-                self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
-            ),
+        expected_output_cols = (
+            self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
         )
+        if isinstance(dataset, DataFrame):
+            expected_output_cols, example_output_pd_df = self._align_expected_output(
+                "fit_predict", dataset, expected_output_cols, output_cols_prefix
+            )
+            output_result, fitted_estimator = model_trainer.train_fit_predict(
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_list=expected_output_cols,
+                example_output_pd_df=example_output_pd_df,
+            )
+        else:
+            output_result, fitted_estimator = model_trainer.train_fit_predict(
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_list=expected_output_cols,
+            )
         self._sklearn_object = fitted_estimator
         self._is_fitted = True
         return output_result
@@ -685,12 +693,41 @@ class IterativeImputer(BaseTransformer):
 
         return rv
 
-    def _align_expected_output_names(
-        self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str
-    ) -> List[str]:
+    def _align_expected_output(
+        self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str,
+    ) -> Tuple[List[str], pd.DataFrame]:
+        """ Run 1 line of data with the desired method, and return one tuple that consists of the output column names 
+        and output dataframe with 1 line.
+        If the method is fit_predict, run 2 lines of data. 
+        """
         # in case the inferred output column names dimension is different 
         # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
-        sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
+
+        # For fit_predict method, a minimum of 2 is required by MinCovDet, BayesianGaussianMixture
+        # so change the minimum of number of rows to 2
+        num_examples = 2
+        statement_params = telemetry.get_function_usage_statement_params(
+            project=_PROJECT,
+            subproject=_SUBPROJECT,
+            function_name=telemetry.get_statement_params_full_func_name(
+                inspect.currentframe(), IterativeImputer.__class__.__name__
+            ),
+            api_calls=[Session.call],
+            custom_tags={"autogen": True} if self._autogenerated else None,
+        )
+        if output_cols_prefix == "fit_predict_":
+            if hasattr(self._sklearn_object, "n_clusters"):
+                # cluster classes such as BisectingKMeansTest requires # of examples >= n_clusters
+                num_examples = self._sklearn_object.n_clusters
+            elif hasattr(self._sklearn_object, "min_samples"):
+                # OPTICS default min_samples 5, which requires at least 5 lines of data
+                num_examples = self._sklearn_object.min_samples
+            elif hasattr(self._sklearn_object, "n_neighbors") and hasattr(self._sklearn_object, "n_samples"):
+                # LocalOutlierFactor expects n_neighbors <= n_samples
+                num_examples = self._sklearn_object.n_neighbors
+            sample_pd_df = dataset.select(self.input_cols).limit(num_examples).to_pandas(statement_params=statement_params)
+        else:
+            sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas(statement_params=statement_params)
 
         # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
         # seen during the fit.
@@ -702,12 +739,14 @@ class IterativeImputer(BaseTransformer):
         output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
         if self.sample_weight_col:
             output_df_columns_set -= set(self.sample_weight_col)
+        
         # if the dimension of inferred output column names is correct; use it
         if len(expected_output_cols_list) == len(output_df_columns_set):
-            return expected_output_cols_list
+            return expected_output_cols_list, output_df_pd
         # otherwise, use the sklearn estimator's output
         else:
-            return sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
+            expected_output_cols_list = sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
+            return expected_output_cols_list, output_df_pd[expected_output_cols_list]
 
     @available_if(original_estimator_has_callable("predict_proba"))  # type: ignore[misc]
     @telemetry.send_api_usage_telemetry(
@@ -753,7 +792,7 @@ class IterativeImputer(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 
@@ -819,7 +858,7 @@ class IterativeImputer(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
         elif isinstance(dataset, pd.DataFrame):
@@ -882,7 +921,7 @@ class IterativeImputer(BaseTransformer):
                 drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )
 
@@ -947,7 +986,7 @@ class IterativeImputer(BaseTransformer):
                 drop_input_cols = self._drop_input_cols,
                 expected_output_cols_type="float",
             )
-            expected_output_cols = self._align_expected_output_names(
+            expected_output_cols, _ = self._align_expected_output(
                 inference_method, dataset, expected_output_cols, output_cols_prefix
             )