snowflake-ml-python 1.6.1__py3-none-any.whl → 1.6.3__py3-none-any.whl

snowflake/ml/modeling/cluster/bisecting_k_means.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
 )
@@ -561,12 +558,23 @@ class BisectingKMeans(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
@@ -591,6 +599,7 @@ class BisectingKMeans(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,
@@ -647,12 +656,41 @@ class BisectingKMeans(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(), BisectingKMeans.__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.
@@ -664,12 +702,14 @@ class BisectingKMeans(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(
@@ -715,7 +755,7 @@ class BisectingKMeans(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
             )
 
@@ -781,7 +821,7 @@ class BisectingKMeans(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):
@@ -844,7 +884,7 @@ class BisectingKMeans(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
             )
 
@@ -909,7 +949,7 @@ class BisectingKMeans(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/cluster/dbscan.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
 )
@@ -525,12 +522,23 @@ class DBSCAN(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
@@ -553,6 +561,7 @@ class DBSCAN(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,
@@ -609,12 +618,41 @@ class DBSCAN(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(), DBSCAN.__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.
@@ -626,12 +664,14 @@ class DBSCAN(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(
@@ -677,7 +717,7 @@ class DBSCAN(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
             )
 
@@ -743,7 +783,7 @@ class DBSCAN(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):
@@ -806,7 +846,7 @@ class DBSCAN(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
             )
 
@@ -871,7 +911,7 @@ class DBSCAN(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/cluster/feature_agglomeration.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
 )
@@ -559,12 +556,23 @@ class FeatureAgglomeration(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
@@ -589,6 +597,7 @@ class FeatureAgglomeration(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,
@@ -645,12 +654,41 @@ class FeatureAgglomeration(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(), FeatureAgglomeration.__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.
@@ -662,12 +700,14 @@ class FeatureAgglomeration(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(
@@ -713,7 +753,7 @@ class FeatureAgglomeration(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
             )
 
@@ -779,7 +819,7 @@ class FeatureAgglomeration(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):
@@ -842,7 +882,7 @@ class FeatureAgglomeration(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
             )
 
@@ -907,7 +947,7 @@ class FeatureAgglomeration(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/cluster/k_means.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
 )
@@ -556,12 +553,23 @@ class KMeans(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
@@ -586,6 +594,7 @@ class KMeans(BaseTransformer):
         """
         self._infer_input_output_cols(dataset)
         super()._check_dataset_type(dataset)
+
         model_trainer = ModelTrainerBuilder.build_fit_transform(
             estimator=self._sklearn_object,
             dataset=dataset,
@@ -642,12 +651,41 @@ class KMeans(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(), KMeans.__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.
@@ -659,12 +697,14 @@ class KMeans(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(
@@ -710,7 +750,7 @@ class KMeans(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
             )
 
@@ -776,7 +816,7 @@ class KMeans(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):
@@ -839,7 +879,7 @@ class KMeans(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
             )
 
@@ -904,7 +944,7 @@ class KMeans(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
             )