snowflake-ml-python 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

snowflake/ml/modeling/impute/knn_imputer.py

@@ -33,6 +33,15 @@ from snowflake.ml.modeling._internal.transformer_protocols import (
     BatchInferenceKwargsTypedDict,
     ScoreKwargsTypedDict
 )
+from snowflake.ml.model._signatures import utils as model_signature_utils
+from snowflake.ml.model.model_signature import (
+    BaseFeatureSpec,
+    DataType,
+    FeatureSpec,
+    ModelSignature,
+    _infer_signature,
+    _rename_signature_with_snowflake_identifiers,
+)
 
 from snowflake.ml.modeling._internal.model_transformer_builder import ModelTransformerBuilder
 
@@ -43,16 +52,6 @@ from snowflake.ml.modeling._internal.estimator_utils import (
     validate_sklearn_args,
 )
 
-from snowflake.ml.model.model_signature import (
-    DataType,
-    FeatureSpec,
-    ModelSignature,
-    _infer_signature,
-    _rename_signature_with_snowflake_identifiers,
-    BaseFeatureSpec,
-)
-from snowflake.ml.model._signatures import utils as model_signature_utils
-
 _PROJECT = "ModelDevelopment"
 # Derive subproject from module name by removing "sklearn"
 # and converting module name from underscore to CamelCase
@@ -61,12 +60,6 @@ _SUBPROJECT = "".join([s.capitalize() for s in "sklearn.impute".replace("sklearn
 
 DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
-def _is_fit_transform_method_enabled() -> Callable[[Any], bool]:	
-    def check(self: BaseTransformer) -> TypeGuard[Callable[..., object]]:	
-        return False and callable(getattr(self._sklearn_object, "fit_transform", None))	
-    return check	
-
-
 class KNNImputer(BaseTransformer):
     r"""Imputation for completing missing values using k-Nearest Neighbors
     For more details on this class, see [sklearn.impute.KNNImputer]
@@ -250,12 +243,7 @@ class KNNImputer(BaseTransformer):
         )
         return selected_cols
 
-    @telemetry.send_api_usage_telemetry(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
-    def fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "KNNImputer":
+    def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "KNNImputer":
         """Fit the imputer on X
         For more details on this function, see [sklearn.impute.KNNImputer.fit]
         (https://scikit-learn.org/stable/modules/generated/sklearn.impute.KNNImputer.html#sklearn.impute.KNNImputer.fit)
@@ -282,12 +270,14 @@ class KNNImputer(BaseTransformer):
 
             self._snowpark_cols = dataset.select(self.input_cols).columns
 
-             # If we are already in a stored procedure, no need to kick off another one.
+            # If we are already in a stored procedure, no need to kick off another one.
             if SNOWML_SPROC_ENV in os.environ:
                 statement_params = telemetry.get_function_usage_statement_params(
                     project=_PROJECT,
                     subproject=_SUBPROJECT,
-                    function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), KNNImputer.__class__.__name__),
+                    function_name=telemetry.get_statement_params_full_func_name(
+                        inspect.currentframe(), KNNImputer.__class__.__name__
+                    ),
                     api_calls=[Session.call],
                     custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
                 )
@@ -308,27 +298,24 @@ class KNNImputer(BaseTransformer):
         )
         self._sklearn_object = model_trainer.train()
         self._is_fitted = True
-        self._get_model_signatures(dataset)
+        self._generate_model_signatures(dataset)
         return self
 
     def _batch_inference_validate_snowpark(
         self,
         dataset: DataFrame,
         inference_method: str,
-    ) -> List[str]:
-        """Util method to run validate that batch inference can be run on a snowpark dataframe and
-        return the available package that exists in the snowflake anaconda channel
+    ) -> None:
+        """Util method to run validate that batch inference can be run on a snowpark dataframe.
 
         Args:
             dataset: snowpark dataframe
             inference_method: the inference method such as predict, score...
-
+            
         Raises:
             SnowflakeMLException: If the estimator is not fitted, raise error
             SnowflakeMLException: If the session is None, raise error
 
-        Returns:
-            A list of available package that exists in the snowflake anaconda channel
         """
         if not self._is_fitted:
             raise exceptions.SnowflakeMLException(
@@ -346,9 +333,7 @@ class KNNImputer(BaseTransformer):
                     "Session must not specified for snowpark dataset."
                 ),
             )
-        # Validate that key package version in user workspace are supported in snowflake conda channel
-        return pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
-            pkg_versions=self._get_dependencies(), session=session, subproject=_SUBPROJECT)
+
 
     @available_if(original_estimator_has_callable("predict"))  # type: ignore[misc]
     @telemetry.send_api_usage_telemetry(
@@ -382,7 +367,9 @@ class KNNImputer(BaseTransformer):
             # when it is classifier, infer the datatype from label columns
             if expected_type_inferred == "" and 'predict' in self.model_signatures:
                 # Batch inference takes a single expected output column type. Use the first columns type for now.
-                label_cols_signatures = [row for row in self.model_signatures['predict'].outputs if row.name in self.output_cols]
+                label_cols_signatures = [
+                    row for row in self.model_signatures['predict'].outputs if row.name in self.output_cols
+                ]
                 if len(label_cols_signatures) == 0:
                     error_str = f"Output columns {self.output_cols} do not match model signatures {self.model_signatures['predict'].outputs}."
                     raise exceptions.SnowflakeMLException(
@@ -390,25 +377,23 @@ class KNNImputer(BaseTransformer):
                         original_exception=ValueError(error_str),
                     )
 
-                expected_type_inferred = convert_sp_to_sf_type(
-                    label_cols_signatures[0].as_snowpark_type()
-                )
+                expected_type_inferred = convert_sp_to_sf_type(label_cols_signatures[0].as_snowpark_type())
             
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
-            assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+            assert isinstance(
+                dataset._session, Session
+            )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
 
             transform_kwargs = dict(
-                session = dataset._session,
-                dependencies = self._deps,
-                drop_input_cols = self._drop_input_cols,
-                expected_output_cols_type = expected_type_inferred,
+                session=dataset._session,
+                dependencies=self._deps,
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_type=expected_type_inferred,
             )
 
         elif isinstance(dataset, pd.DataFrame):
-            transform_kwargs = dict(
-                snowpark_input_cols = self._snowpark_cols,
-                drop_input_cols = self._drop_input_cols
-                )
+            transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
 
         transform_handlers = ModelTransformerBuilder.build(
             dataset=dataset,
@@ -450,7 +435,7 @@ class KNNImputer(BaseTransformer):
             Transformed dataset.
         """
         super()._check_dataset_type(dataset)
-        inference_method="transform"
+        inference_method = "transform"
 
         # This dictionary contains optional kwargs for batch inference. These kwargs
         # are specific to the type of dataset used. 
@@ -480,24 +465,19 @@ class KNNImputer(BaseTransformer):
                     if all(x == output_types[0] for x in output_types) and len(output_types) == len(self.output_cols):
                         expected_dtype = convert_sp_to_sf_type(output_types[0])
             
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-            )
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
 
             transform_kwargs = dict(
-                session = dataset._session,
-                dependencies = self._deps,
-                drop_input_cols = self._drop_input_cols,
-                expected_output_cols_type = expected_dtype,
+                session=dataset._session,
+                dependencies=self._deps,
+                drop_input_cols=self._drop_input_cols,
+                expected_output_cols_type=expected_dtype,
             )
 
         elif isinstance(dataset, pd.DataFrame):
-            transform_kwargs = dict(
-                snowpark_input_cols = self._snowpark_cols,
-                drop_input_cols = self._drop_input_cols
-                )
+            transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
 
         transform_handlers = ModelTransformerBuilder.build(
             dataset=dataset,
@@ -516,7 +496,11 @@ class KNNImputer(BaseTransformer):
         return output_df
     
     @available_if(original_estimator_has_callable("fit_predict"))  # type: ignore[misc]
-    def fit_predict(self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "fit_predict_",) -> Union[DataFrame, pd.DataFrame]:
+    def fit_predict(
+        self,
+        dataset: Union[DataFrame, pd.DataFrame],
+        output_cols_prefix: str = "fit_predict_",
+    ) -> Union[DataFrame, pd.DataFrame]:
         """ Method not supported for this class.
 
 
@@ -541,22 +525,106 @@ class KNNImputer(BaseTransformer):
         )
         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_list=(
+                self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
+            ),
         )
         self._sklearn_object = fitted_estimator
         self._is_fitted = True
         return output_result
 
+            
+    @available_if(original_estimator_has_callable("fit_transform"))  # type: ignore[misc]
+    def fit_transform(self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "fit_transform_",) -> Union[DataFrame, pd.DataFrame]:
+        """ Fit to data, then transform it
+        For more details on this function, see [sklearn.impute.KNNImputer.fit_transform]
+        (https://scikit-learn.org/stable/modules/generated/sklearn.impute.KNNImputer.html#sklearn.impute.KNNImputer.fit_transform)
+
+
+        Raises:
+            TypeError: Supported dataset types: snowpark.DataFrame, pandas.DataFrame.
 
-    @available_if(_is_fit_transform_method_enabled())  # type: ignore[misc]
-    def fit_transform(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[Any, npt.NDArray[Any]]:
-        """ 
+        Args:
+            dataset: Union[snowflake.snowpark.DataFrame, pandas.DataFrame]
+                Snowpark or Pandas DataFrame.
+        output_cols_prefix: Prefix for the response columns
         Returns:
             Transformed dataset.
         """
-        self.fit(dataset)
-        assert self._sklearn_object is not None
-        return self._sklearn_object.embedding_
+        self._infer_input_output_cols(dataset)
+        super()._check_dataset_type(dataset)
+        model_trainer = ModelTrainerBuilder.build_fit_transform(
+            estimator=self._sklearn_object,
+            dataset=dataset,
+            input_cols=self.input_cols,
+            label_cols=self.label_cols,
+            sample_weight_col=self.sample_weight_col,
+            autogenerated=self._autogenerated,
+            subproject=_SUBPROJECT,
+        )
+        output_result, fitted_estimator = model_trainer.train_fit_transform(
+            drop_input_cols=self._drop_input_cols,
+            expected_output_cols_list=self.output_cols,
+        )
+        self._sklearn_object = fitted_estimator
+        self._is_fitted = True
+        return output_result
+
+
+    def _get_output_column_names(self, output_cols_prefix: str, output_cols: Optional[List[str]] = None) -> List[str]:
+        """ Returns the list of output columns for predict_proba(), decision_function(), etc.. functions.
+        Returns a list with output_cols_prefix as the only element if the estimator is not a classifier.
+        """
+        output_cols_prefix = identifier.resolve_identifier(output_cols_prefix)
+        # The following condition is introduced for kneighbors methods, and not used in other methods
+        if output_cols:
+            output_cols = [
+                identifier.concat_names([output_cols_prefix, identifier.resolve_identifier(c)])
+                for c in output_cols
+            ]
+        elif getattr(self._sklearn_object, "classes_", None) is None:
+            output_cols = [output_cols_prefix]
+        elif self._sklearn_object is not None:
+            classes = self._sklearn_object.classes_
+            if isinstance(classes, numpy.ndarray):
+                output_cols = [f'{output_cols_prefix}{str(c)}' for c in classes.tolist()]
+            elif isinstance(classes, list) and len(classes) > 0 and isinstance(classes[0], numpy.ndarray):
+                # If the estimator is a multioutput estimator, classes_ will be a list of ndarrays.
+                output_cols = []
+                for i, cl in enumerate(classes):
+                    # For binary classification, there is only one output column for each class
+                    # ndarray as the two classes are complementary.
+                    if len(cl) == 2:
+                        output_cols.append(f'{output_cols_prefix}{i}_{cl[0]}')
+                    else:
+                        output_cols.extend([
+                            f'{output_cols_prefix}{i}_{c}' for c in cl.tolist()
+                        ])
+        else:
+            output_cols = []
+
+        # Make sure column names are valid snowflake identifiers.
+        assert output_cols is not None  # Make MyPy happy
+        rv = [identifier.rename_to_valid_snowflake_identifier(c) for c in output_cols]
+
+        return rv
+
+    def _align_expected_output_names(
+        self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str
+    ) -> List[str]:
+        # 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
+        output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
+        output_df_columns = list(output_df_pd.columns)
+        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
+        # otherwise, use the sklearn estimator's output
+        else:
+            return sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
 
     @available_if(original_estimator_has_callable("predict_proba"))  # type: ignore[misc]
     @telemetry.send_api_usage_telemetry(
@@ -588,24 +656,26 @@ class KNNImputer(BaseTransformer):
         # are specific to the type of dataset used. 
         transform_kwargs: BatchInferenceKwargsTypedDict = dict() 
 
+        expected_output_cols = self._get_output_column_names(output_cols_prefix)
+
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-            )
-            assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+            assert isinstance(
+                dataset._session, Session
+            )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
             transform_kwargs = dict(
                 session=dataset._session,
                 dependencies=self._deps,
-                drop_input_cols = self._drop_input_cols,
+                drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
+            expected_output_cols = self._align_expected_output_names(
+                inference_method, dataset, expected_output_cols, output_cols_prefix
+            )
 
         elif isinstance(dataset, pd.DataFrame):
-            transform_kwargs = dict(
-                snowpark_input_cols = self._snowpark_cols,
-                drop_input_cols = self._drop_input_cols
-                )
+            transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
 
         transform_handlers = ModelTransformerBuilder.build(
             dataset=dataset,
@@ -617,7 +687,7 @@ class KNNImputer(BaseTransformer):
         output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
             inference_method=inference_method,
             input_cols=self.input_cols,
-            expected_output_cols=self._get_output_column_names(output_cols_prefix),
+            expected_output_cols=expected_output_cols,
             **transform_kwargs
         )
         return output_df
@@ -647,29 +717,30 @@ class KNNImputer(BaseTransformer):
             Output dataset with log probability of the sample for each class in the model.
         """
         super()._check_dataset_type(dataset)
-        inference_method="predict_log_proba"
+        inference_method = "predict_log_proba"
+        expected_output_cols = self._get_output_column_names(output_cols_prefix)
 
         # This dictionary contains optional kwargs for batch inference. These kwargs
         # are specific to the type of dataset used. 
         transform_kwargs: BatchInferenceKwargsTypedDict = dict()  
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-            )
-            assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+            assert isinstance(
+                dataset._session, Session
+            )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
             transform_kwargs = dict(
                 session=dataset._session,
                 dependencies=self._deps,
-                drop_input_cols = self._drop_input_cols,
+                drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
+            expected_output_cols = self._align_expected_output_names(
+                inference_method, dataset, expected_output_cols, output_cols_prefix
+            )
         elif isinstance(dataset, pd.DataFrame):
-            transform_kwargs = dict(
-                snowpark_input_cols = self._snowpark_cols,
-                drop_input_cols = self._drop_input_cols
-                )
+            transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
 
         transform_handlers = ModelTransformerBuilder.build(
             dataset=dataset,
@@ -682,7 +753,7 @@ class KNNImputer(BaseTransformer):
         output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
             inference_method=inference_method,
             input_cols=self.input_cols,
-            expected_output_cols=self._get_output_column_names(output_cols_prefix),
+            expected_output_cols=expected_output_cols,
             **transform_kwargs
         )
         return output_df
@@ -708,30 +779,32 @@ class KNNImputer(BaseTransformer):
             Output dataset with results of the decision function for the samples in input dataset.
         """
         super()._check_dataset_type(dataset)
-        inference_method="decision_function"
+        inference_method = "decision_function"
 
         # This dictionary contains optional kwargs for batch inference. These kwargs
         # are specific to the type of dataset used. 
         transform_kwargs: BatchInferenceKwargsTypedDict = dict()  
 
+        expected_output_cols = self._get_output_column_names(output_cols_prefix)
+
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-            )
-            assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+            assert isinstance(
+                dataset._session, Session
+            )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
             transform_kwargs = dict(
                 session=dataset._session,
                 dependencies=self._deps,
-                drop_input_cols = self._drop_input_cols,
+                drop_input_cols=self._drop_input_cols,
                 expected_output_cols_type="float",
             )
+            expected_output_cols = self._align_expected_output_names(
+                inference_method, dataset, expected_output_cols, output_cols_prefix
+            )
 
         elif isinstance(dataset, pd.DataFrame):
-            transform_kwargs = dict(
-                snowpark_input_cols = self._snowpark_cols,
-                drop_input_cols = self._drop_input_cols
-                )
+            transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
 
         transform_handlers = ModelTransformerBuilder.build(
             dataset=dataset,
@@ -744,7 +817,7 @@ class KNNImputer(BaseTransformer):
         output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
             inference_method=inference_method,
             input_cols=self.input_cols,
-            expected_output_cols=self._get_output_column_names(output_cols_prefix),
+            expected_output_cols=expected_output_cols,
             **transform_kwargs
         )
         return output_df
@@ -773,17 +846,17 @@ class KNNImputer(BaseTransformer):
             Output dataset with probability of the sample for each class in the model.
         """
         super()._check_dataset_type(dataset)
-        inference_method="score_samples"
+        inference_method = "score_samples"
 
         # This dictionary contains optional kwargs for batch inference. These kwargs
         # are specific to the type of dataset used. 
         transform_kwargs: BatchInferenceKwargsTypedDict = dict()
 
+        expected_output_cols = self._get_output_column_names(output_cols_prefix)
+
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-            )
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
             transform_kwargs = dict(
                 session=dataset._session,
@@ -791,6 +864,9 @@ class KNNImputer(BaseTransformer):
                 drop_input_cols = self._drop_input_cols,
                 expected_output_cols_type="float",
             )
+            expected_output_cols = self._align_expected_output_names(
+                inference_method, dataset, expected_output_cols, output_cols_prefix
+            )
 
         elif isinstance(dataset, pd.DataFrame):
             transform_kwargs = dict(
@@ -809,7 +885,7 @@ class KNNImputer(BaseTransformer):
         output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
             inference_method=inference_method,
             input_cols=self.input_cols,
-            expected_output_cols=self._get_output_column_names(output_cols_prefix),
+            expected_output_cols=expected_output_cols,
             **transform_kwargs
         )
         return output_df
@@ -842,17 +918,15 @@ class KNNImputer(BaseTransformer):
         transform_kwargs: ScoreKwargsTypedDict = dict()  
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method="score",
-            )
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method="score")
+            self._deps = self._get_dependencies()
             selected_cols = self._get_active_columns()
             if len(selected_cols) > 0:
                 dataset = dataset.select(selected_cols)
             assert isinstance(dataset._session, Session) # keep mypy happy
             transform_kwargs = dict(
                 session=dataset._session,
-                dependencies=["snowflake-snowpark-python"] + self._deps,
+                dependencies=self._deps,
                 score_sproc_imports=['sklearn'],
             )
         elif isinstance(dataset, pd.DataFrame):
@@ -917,11 +991,8 @@ class KNNImputer(BaseTransformer):
 
         if isinstance(dataset, DataFrame):
 
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-
-            )
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
             transform_kwargs = dict(
                 session = dataset._session,
@@ -954,50 +1025,84 @@ class KNNImputer(BaseTransformer):
         )
         return output_df
 
+    
+
+    def to_sklearn(self) -> Any:
+        """Get sklearn.impute.KNNImputer object.
+        """
+        if self._sklearn_object is None:
+            self._sklearn_object = self._create_sklearn_object()
+        return self._sklearn_object
+
+    def to_xgboost(self) -> Any:
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.METHOD_NOT_ALLOWED,
+            original_exception=AttributeError(
+                modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
+                    "to_xgboost()", 
+                    "to_sklearn()"
+                )
+            ),
+        )
 
-    def _get_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
+    def to_lightgbm(self) -> Any:
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.METHOD_NOT_ALLOWED,
+            original_exception=AttributeError(
+                modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
+                    "to_lightgbm()", 
+                    "to_sklearn()"
+                )
+            ),
+        )
+
+    def _get_dependencies(self) -> List[str]:
+        return self._deps
+
+
+    def _generate_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
         self._model_signature_dict = dict()
 
         PROB_FUNCTIONS = ["predict_log_proba", "predict_proba", "decision_function"]
 
-        inputs = list(_infer_signature(dataset[self.input_cols], "input"))
+        inputs = list(_infer_signature(dataset[self.input_cols], "input", use_snowflake_identifiers=True))
         outputs: List[BaseFeatureSpec] = []
         if hasattr(self, "predict"):
             # keep mypy happy
-            assert self._sklearn_object is not None and hasattr(self._sklearn_object, "_estimator_type") 
+            assert self._sklearn_object is not None and hasattr(self._sklearn_object, "_estimator_type")
             # For classifier, the type of predict is the same as the type of label
-            if self._sklearn_object._estimator_type == 'classifier':
-                 # label columns is the desired type for output
+            if self._sklearn_object._estimator_type == "classifier":
+                # label columns is the desired type for output
                 outputs = list(_infer_signature(dataset[self.label_cols], "output", use_snowflake_identifiers=True))
                 # rename the output columns
                 outputs = list(model_signature_utils.rename_features(outputs, self.output_cols))
-                self._model_signature_dict["predict"] = ModelSignature(inputs,
-                                                                       ([] if self._drop_input_cols else inputs)
-                                                                       + outputs)
+                self._model_signature_dict["predict"] = ModelSignature(
+                    inputs, ([] if self._drop_input_cols else inputs) + outputs
+                )
             # For mixture models that use the density mixin, `predict` returns the argmax of the log prob.
             # For outlier models, returns -1 for outliers and 1 for inliers.
-            # Clusterer returns int64 cluster labels. 
+            # Clusterer returns int64 cluster labels.
             elif self._sklearn_object._estimator_type in ["DensityEstimator", "clusterer", "outlier_detector"]:
                 outputs = [FeatureSpec(dtype=DataType.INT64, name=c) for c in self.output_cols]
-                self._model_signature_dict["predict"] = ModelSignature(inputs,
-                                                                       ([] if self._drop_input_cols else inputs)
-                                                                       + outputs)
-            
+                self._model_signature_dict["predict"] = ModelSignature(
+                    inputs, ([] if self._drop_input_cols else inputs) + outputs
+                )
+
             # For regressor, the type of predict is float64
-            elif self._sklearn_object._estimator_type == 'regressor':
+            elif self._sklearn_object._estimator_type == "regressor":
                 outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in self.output_cols]
-                self._model_signature_dict["predict"] = ModelSignature(inputs,
-                                                                       ([] if self._drop_input_cols else inputs)
-                                                                       + outputs)
-                
+                self._model_signature_dict["predict"] = ModelSignature(
+                    inputs, ([] if self._drop_input_cols else inputs) + outputs
+                )
+
         for prob_func in PROB_FUNCTIONS:
             if hasattr(self, prob_func):
                 output_cols_prefix: str = f"{prob_func}_"
                 output_column_names = self._get_output_column_names(output_cols_prefix)
                 outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in output_column_names]
-                self._model_signature_dict[prob_func] = ModelSignature(inputs,
-                                                                       ([] if self._drop_input_cols else inputs)
-                                                                       + outputs)
+                self._model_signature_dict[prob_func] = ModelSignature(
+                    inputs, ([] if self._drop_input_cols else inputs) + outputs
+                )
 
         # Output signature names may still need to be renamed, since they were not created with `_infer_signature`.
         items = list(self._model_signature_dict.items())
@@ -1010,10 +1115,10 @@ class KNNImputer(BaseTransformer):
         """Returns model signature of current class.
 
         Raises:
-            exceptions.SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
+            SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
 
         Returns:
-            Dict[str, ModelSignature]: each method and its input output signature
+            Dict with each method and its input output signature
         """
         if self._model_signature_dict is None:
             raise exceptions.SnowflakeMLException(
@@ -1021,35 +1126,3 @@ class KNNImputer(BaseTransformer):
                 original_exception=RuntimeError("Estimator not fitted before accessing property model_signatures!"),
             )
         return self._model_signature_dict
-
-    def to_sklearn(self) -> Any:
-        """Get sklearn.impute.KNNImputer object.
-        """
-        if self._sklearn_object is None:
-            self._sklearn_object = self._create_sklearn_object()
-        return self._sklearn_object
-
-    def to_xgboost(self) -> Any:
-        raise exceptions.SnowflakeMLException(
-            error_code=error_codes.METHOD_NOT_ALLOWED,
-            original_exception=AttributeError(
-                modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
-                    "to_xgboost()", 
-                    "to_sklearn()"
-                )
-            ),
-        )
-
-    def to_lightgbm(self) -> Any:
-        raise exceptions.SnowflakeMLException(
-            error_code=error_codes.METHOD_NOT_ALLOWED,
-            original_exception=AttributeError(
-                modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
-                    "to_lightgbm()", 
-                    "to_sklearn()"
-                )
-            ),
-        )
-
-    def _get_dependencies(self) -> List[str]:
-        return self._deps