snowflake-ml-python 1.3.1__py3-none-any.whl → 1.4.1__py3-none-any.whl

snowflake/ml/modeling/decomposition/pca.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
@@ -286,12 +285,7 @@ class PCA(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]) -> "PCA":
+    def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "PCA":
         """Fit the model with X
         For more details on this function, see [sklearn.decomposition.PCA.fit]
         (https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.PCA.html#sklearn.decomposition.PCA.fit)
@@ -318,12 +312,14 @@ class PCA(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(), PCA.__class__.__name__),
+                    function_name=telemetry.get_statement_params_full_func_name(
+                        inspect.currentframe(), PCA.__class__.__name__
+                    ),
                     api_calls=[Session.call],
                     custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
                 )
@@ -344,7 +340,7 @@ class PCA(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(
@@ -418,7 +414,9 @@ class PCA(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(
@@ -426,25 +424,22 @@ class PCA(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()
+            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,
@@ -486,7 +481,7 @@ class PCA(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. 
@@ -523,17 +518,14 @@ class PCA(BaseTransformer):
             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,
@@ -552,7 +544,11 @@ class PCA(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.
 
 
@@ -577,7 +573,9 @@ class PCA(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
@@ -594,6 +592,62 @@ class PCA(BaseTransformer):
         assert self._sklearn_object is not None
         return self._sklearn_object.embedding_
 
+
+    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(
         project=_PROJECT,
@@ -624,24 +678,28 @@ class PCA(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()
+            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,
@@ -653,7 +711,7 @@ class PCA(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
@@ -683,7 +741,8 @@ class PCA(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. 
@@ -694,18 +753,20 @@ class PCA(BaseTransformer):
                 dataset=dataset,
                 inference_method=inference_method,
             )
-            assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
+            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,
@@ -718,7 +779,7 @@ class PCA(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
@@ -744,30 +805,34 @@ class PCA(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()
+            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,
@@ -780,7 +845,7 @@ class PCA(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
@@ -811,12 +876,14 @@ class PCA(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,
@@ -829,6 +896,9 @@ class PCA(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(
@@ -847,7 +917,7 @@ class PCA(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
@@ -994,50 +1064,84 @@ class PCA(BaseTransformer):
         )
         return output_df
 
+    
+
+    def to_sklearn(self) -> Any:
+        """Get sklearn.decomposition.PCA 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_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
+    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())
@@ -1050,10 +1154,10 @@ class PCA(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(
@@ -1061,35 +1165,3 @@ class PCA(BaseTransformer):
                 original_exception=RuntimeError("Estimator not fitted before accessing property model_signatures!"),
             )
         return self._model_signature_dict
-
-    def to_sklearn(self) -> Any:
-        """Get sklearn.decomposition.PCA 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