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

snowflake/ml/modeling/preprocessing/robust_scaler.py

@@ -144,11 +144,7 @@ class RobustScaler(base.BaseTransformer):
     def scale_(self) -> Optional[Dict[str, float]]:
         return None if (not self.with_scaling or not self._state_is_set) else self._scale
 
-    @telemetry.send_api_usage_telemetry(
-        project=base.PROJECT,
-        subproject=base.SUBPROJECT,
-    )
-    def fit(self, dataset: Union[snowpark.DataFrame, pd.DataFrame]) -> "RobustScaler":
+    def _fit(self, dataset: Union[snowpark.DataFrame, pd.DataFrame]) -> "RobustScaler":
         """
         Compute center, scale and quantile values of the dataset.
 

snowflake/ml/modeling/preprocessing/standard_scaler.py

@@ -132,11 +132,7 @@ class StandardScaler(base.BaseTransformer):
         if hasattr(self, "var_"):
             self.var_ = {} if self.with_std else None
 
-    @telemetry.send_api_usage_telemetry(
-        project=base.PROJECT,
-        subproject=base.SUBPROJECT,
-    )
-    def fit(self, dataset: Union[snowpark.DataFrame, pd.DataFrame]) -> "StandardScaler":
+    def _fit(self, dataset: Union[snowpark.DataFrame, pd.DataFrame]) -> "StandardScaler":
         """
         Compute mean and std values of the dataset.
 
@@ -165,11 +161,11 @@ class StandardScaler(base.BaseTransformer):
 
         for i, input_col in enumerate(self.input_cols):
             if self.mean_ is not None:
-                self.mean_[input_col] = float(sklearn_scaler.mean_[i])
+                self.mean_[input_col] = _utils.to_float_if_valid(sklearn_scaler.mean_[i], input_col, "mean_")
             if self.scale_ is not None:
-                self.scale_[input_col] = float(sklearn_scaler.scale_[i])
+                self.scale_[input_col] = _utils.to_float_if_valid(sklearn_scaler.scale_[i], input_col, "scale_")
             if self.var_ is not None:
-                self.var_[input_col] = float(sklearn_scaler.var_[i])
+                self.var_[input_col] = _utils.to_float_if_valid(sklearn_scaler.var_[i], input_col, "var_")
 
     def _fit_snowpark(self, dataset: snowpark.DataFrame) -> None:
         computed_states = self._compute(dataset, self.input_cols, self.custom_states)
@@ -179,14 +175,18 @@ class StandardScaler(base.BaseTransformer):
             numeric_stats = computed_states[input_col]
 
             if self.mean_ is not None:
-                self.mean_[input_col] = float(numeric_stats[_utils.NumericStatistics.MEAN])
+                self.mean_[input_col] = _utils.to_float_if_valid(
+                    numeric_stats[_utils.NumericStatistics.MEAN], input_col, "mean_"
+                )
 
             if self.var_ is not None:
-                self.var_[input_col] = float(numeric_stats[_utils.NumericStatistics.VAR_POP])
+                self.var_[input_col] = _utils.to_float_if_valid(
+                    numeric_stats[_utils.NumericStatistics.VAR_POP], input_col, "var_"
+                )
 
             if self.scale_ is not None:
                 self.scale_[input_col] = sklearn_preprocessing_data._handle_zeros_in_scale(
-                    float(numeric_stats[_utils.NumericStatistics.STDDEV_POP])
+                    _utils.to_float_if_valid(numeric_stats[_utils.NumericStatistics.STDDEV_POP], input_col, "scale_")
                 )
 
     @telemetry.send_api_usage_telemetry(