Control the types on sklearn internal read table

We do this only for KhiopsClassifier and KhiopsRegressor: It is critical
for KhiopsClassifier as it accepts many target types and it is trivial
in the case of KhiopsRegressor.

For KhiopsEncoder and KhiopsCoclustering is less critical and for the
first one it is very complex. We left them as TODO's.

Additionaly, we now also check in the "output type" tests that the
result of predict is correct. Before we only checked only that the
classes_ attribute was ok. This is to further ensure correctness.

Author: folmos-at-orange

khiops/sklearn/estimators.py

@@ -231,7 +231,7 @@ class KhiopsEstimator(ABC, BaseEstimator):
         The name of the column to be used as key.
         **Deprecated** will be removed in Khiops 11.
     internal_sort : bool, optional
-        *Advanced.*: See concrete estimator classes for information about this
+        *Advanced*: See concrete estimator classes for information about this
         parameter.
         **Deprecated** will be removed in Khiops 11. Use the ``auto_sort``
         estimator parameter instead.
@@ -466,7 +466,7 @@ def _transform(
         self,
         ds,
         computation_dir,
-        _transform_create_deployment_model_fun,
+        _transform_prepare_deployment_fun,
         drop_key,
         transformed_file_name,
     ):
@@ -478,11 +478,13 @@ def _transform(
         self._transform_check_dataset(ds)
 
         # Create a deployment dataset
-        # Note: The input dataset is not necessarily ready to be deployed
+        # Note: The input dataset isn't ready for deployment in the case of coclustering
         deployment_ds = self._transform_create_deployment_dataset(ds, computation_dir)
 
-        # Create a deployment dictionary
-        deployment_dictionary_domain = _transform_create_deployment_model_fun(ds)
+        # Create a deployment dictionary and the internal table column dtypes
+        deployment_dictionary_domain, internal_table_column_dtypes = (
+            _transform_prepare_deployment_fun(ds)
+        )
 
         # Deploy the model
         output_table_path = self._transform_deploy_model(
@@ -493,10 +495,36 @@ def _transform(
             transformed_file_name,
         )
 
-        # Post-process to return the correct output type
-        return self._transform_deployment_post_process(
-            deployment_ds, output_table_path, drop_key
-        )
+        # Post-process to return the correct output type and order
+        if deployment_ds.is_in_memory:
+            # Load the table as a dataframe
+            with io.BytesIO(fs.read(output_table_path)) as output_table_stream:
+                output_table_df = read_internal_data_table(
+                    output_table_stream, column_dtypes=internal_table_column_dtypes
+                )
+
+            # On multi-table:
+            # - Reorder the table to the original table order
+            #     - Because transformed data table file is sorted by key
+            # - Drop the key columns if specified
+            if deployment_ds.is_multitable:
+                key_df = deployment_ds.main_table.data_source[
+                    deployment_ds.main_table.key
+                ]
+                output_table_df_or_path = key_df.merge(
+                    output_table_df, on=deployment_ds.main_table.key
+                )
+                if drop_key:
+                    output_table_df_or_path.drop(
+                        deployment_ds.main_table.key, axis=1, inplace=True
+                    )
+            # On mono-table: Return the read dataframe as-is
+            else:
+                output_table_df_or_path = output_table_df
+        else:
+            output_table_df_or_path = output_table_path
+
+        return output_table_df_or_path
 
     def _transform_create_deployment_dataset(self, ds, _):
         """Creates if necessary a new dataset to execute the model deployment
@@ -605,44 +633,6 @@ def _transform_check_dataset(self, ds):
         if ds.table_type == FileTable and self.output_dir is None:
             raise ValueError("'output_dir' is not set but dataset is file-based")
 
-    def _transform_deployment_post_process(
-        self, deployment_ds, output_table_path, drop_key
-    ):
-        # Return a dataframe for dataframe based datasets
-        if deployment_ds.is_in_memory:
-            # Read the transformed table with the internal table settings
-            with io.BytesIO(fs.read(output_table_path)) as output_table_stream:
-                output_table_df = read_internal_data_table(output_table_stream)
-
-            # On multi-table:
-            # - Reorder the table to the original table order
-            #     - Because transformed data table file is sorted by key
-            # - Drop the key columns if specified
-            if deployment_ds.is_multitable:
-                key_df = deployment_ds.main_table.data_source[
-                    deployment_ds.main_table.key
-                ]
-                output_table_df_or_path = key_df.merge(
-                    output_table_df, on=deployment_ds.main_table.key
-                )
-                if drop_key:
-                    output_table_df_or_path.drop(
-                        deployment_ds.main_table.key, axis=1, inplace=True
-                    )
-            # On mono-table: Return the read dataframe as-is
-            else:
-                output_table_df_or_path = output_table_df
-        # Return a file path for file based datasets
-        else:
-            output_table_df_or_path = output_table_path
-
-        assert isinstance(
-            output_table_df_or_path, (str, pd.DataFrame)
-        ), type_error_message(
-            "output_table_df_or_path", output_table_df_or_path, str, pd.DataFrame
-        )
-        return output_table_df_or_path
-
     def _create_computation_dir(self, method_name):
         """Creates a temporary computation directory"""
         return kh.get_runner().create_temp_dir(
@@ -1260,7 +1250,7 @@ def predict(self, X):
             y_pred = super()._transform(
                 ds,
                 computation_dir,
-                self._transform_prepare_deployment_model_for_predict,
+                self._transform_prepare_deployment_for_predict,
                 False,
                 "predict.txt",
             )
@@ -1366,16 +1356,12 @@ def _transform_create_deployment_dataset(self, ds, computation_dir):
 
         return Dataset(deploy_dataset_spec)
 
-    def _transform_prepare_deployment_model_for_predict(self, _):
-        return self.model_.copy()
-
-    def _transform_deployment_post_process(
-        self, deployment_ds, output_table_path, drop_key
-    ):
-        assert deployment_ds.is_multitable
-        return super()._transform_deployment_post_process(
-            deployment_ds, output_table_path, drop_key
-        )
+    def _transform_prepare_deployment_for_predict(self, _):
+        # TODO
+        #  Replace the second return value (the output columns' dtypes) with a proper
+        #  value instead of `None`. In the current state, it will use pandas type
+        #  auto-detection to load the internal table into memory.
+        return self.model_.copy(), None
 
     def fit_predict(self, X, y=None, **kwargs):
         """Performs clustering on X and returns result (instead of labels)"""
@@ -1412,6 +1398,7 @@ def __init__(
         self.specific_pairs = specific_pairs
         self.all_possible_pairs = all_possible_pairs
         self.construction_rules = construction_rules
+        self._original_target_dtype = None
         self._predicted_target_meta_data_tag = None
 
         # Deprecation message for 'key' constructor parameter
@@ -1619,6 +1606,22 @@ def _fit_training_post_process(self, ds):
         # Call parent method
         super()._fit_training_post_process(ds)
 
+        # Save the target and key column dtype's
+        if ds.is_in_memory:
+            if self._original_target_dtype is None:
+                self._original_target_dtype = ds.target_column.dtype
+            if ds.main_table.key is not None:
+                self._original_key_dtypes = {}
+                for column_id in ds.main_table.key:
+                    self._original_key_dtypes[column_id] = (
+                        ds.main_table.get_column_dtype(column_id)
+                    )
+            else:
+                self._original_key_dtypes = None
+        else:
+            self._original_target_dtype = None
+            self._original_key_dtypes = None
+
         # Set the target variable name
         self.model_target_variable_name_ = get_khiops_variable_name(ds.target_column_id)
 
@@ -1794,6 +1797,7 @@ def __init__(
         )
         # Data to be specified by inherited classes
         self._predicted_target_meta_data_tag = None
+        self._predicted_target_name_prefix = None
         self.n_evaluated_features = n_evaluated_features
         self.n_selected_features = n_selected_features
 
@@ -1821,7 +1825,7 @@ def predict(self, X):
             y_pred = super()._transform(
                 ds,
                 computation_dir,
-                self._transform_prepare_deployment_model_for_predict,
+                self._transform_prepare_deployment_for_predict,
                 True,
                 "predict.txt",
             )
@@ -1849,7 +1853,7 @@ def _fit_prepare_training_function_inputs(self, ds, computation_dir):
 
         return args, kwargs
 
-    def _transform_prepare_deployment_model_for_predict(self, ds):
+    def _transform_prepare_deployment_for_predict(self, ds):
         assert (
             self._predicted_target_meta_data_tag is not None
         ), "Predicted target metadata tag is not set"
@@ -1874,7 +1878,20 @@ def _transform_prepare_deployment_model_for_predict(self, ds):
         if self.model_target_variable_name_ not in list(ds.main_table.column_ids):
             model_dictionary.remove_variable(self.model_target_variable_name_)
 
-        return model_copy
+        # Create the output column dtype dict
+        if ds.is_in_memory:
+            predicted_target_column_name = (
+                self._predicted_target_name_prefix + self.model_target_variable_name_
+            )
+            output_columns_dtype = {
+                predicted_target_column_name: self._original_target_dtype
+            }
+            if self.is_multitable_model_:
+                output_columns_dtype.update(self._original_key_dtypes)
+        else:
+            output_columns_dtype = None
+
+        return model_copy, output_columns_dtype
 
     def get_feature_used_statistics(self, modeling_report):
         # Extract, from the modeling report, names, levels, weights and importances
@@ -1889,7 +1906,7 @@ def get_feature_used_statistics(self, modeling_report):
                     for var in modeling_report.selected_variables
                 ]
             )
-        # Return empty arrays if not selected_variables is available
+        # Return empty arrays if no selected variables are available
         else:
             feature_used_names_ = np.array([], dtype=np.dtype("<U1"))
             feature_used_importances_ = np.array([])
@@ -2075,6 +2092,7 @@ def __init__(
         self.group_target_value = group_target_value
         self._khiops_model_prefix = "SNB_"
         self._predicted_target_meta_data_tag = "Prediction"
+        self._predicted_target_name_prefix = "Predicted"
 
     def __sklearn_tags__(self):
         # If we don't implement this trivial method it's not found by the sklearn. This
@@ -2183,12 +2201,6 @@ def _fit_training_post_process(self, ds):
         # Call the parent's method
         super()._fit_training_post_process(ds)
 
-        # Save the target datatype
-        if ds.is_in_memory:
-            self._original_target_dtype = ds.target_column.dtype
-        else:
-            self._original_target_dtype = None
-
         # Save class values in the order of deployment
         self.classes_ = []
         for variable in self._get_main_dictionary().variables:
@@ -2260,37 +2272,15 @@ def predict(self, X):
         # Call the parent's method
         y_pred = super().predict(X)
 
-        # Adjust the data type according to the original target type
-        # Note: String is coerced explictly because astype does not work as expected
+        # Convert to numpy if it is in memory
         if isinstance(y_pred, pd.DataFrame):
-            # Transform to numpy.ndarray
-            y_pred = y_pred.to_numpy(copy=False).ravel()
-
-            # If integer and string just transform
-            if pd.api.types.is_integer_dtype(self._original_target_dtype):
-                y_pred = y_pred.astype(self._original_target_dtype)
-            # If str transform to str
-            # Note: If the original type is None then it was learned with a file dataset
-            elif self._original_target_dtype is None or pd.api.types.is_string_dtype(
-                self._original_target_dtype
-            ):
-                y_pred = y_pred.astype(str, copy=False)
-            # If category first coerce the type to the categories' type
-            else:
-                assert isinstance(self._original_target_dtype, pd.CategoricalDtype), (
-                    "_original_target_dtype is not categorical"
-                    f", it is '{self._original_target_dtype}'"
-                )
-                if pd.api.types.is_integer_dtype(
-                    self._original_target_dtype.categories.dtype
-                ):
-                    y_pred = y_pred.astype(
-                        self._original_target_dtype.categories.dtype, copy=False
-                    )
-                else:
-                    y_pred = y_pred.astype(str, copy=False)
+            y_pred = y_pred[
+                self._predicted_target_name_prefix + self.model_target_variable_name_
+            ].to_numpy(copy=False)
 
-        assert isinstance(y_pred, (str, np.ndarray)), "Expected str or np.array"
+        assert isinstance(y_pred, (np.ndarray, str)), type_error_message(
+            "y_pred", y_pred, np.ndarray, str
+        )
         return y_pred
 
     def predict_proba(self, X):
@@ -2336,7 +2326,7 @@ def predict_proba(self, X):
             y_probas = self._transform(
                 ds,
                 computation_dir,
-                self._transform_prepare_deployment_model_for_predict_proba,
+                self._transform_prepare_deployment_for_predict_proba,
                 True,
                 "predict_proba.txt",
             )
@@ -2359,7 +2349,7 @@ def predict_proba(self, X):
         assert isinstance(y_probas, (str, np.ndarray)), "Expected str or np.ndarray"
         return y_probas
 
-    def _transform_prepare_deployment_model_for_predict_proba(self, ds):
+    def _transform_prepare_deployment_for_predict_proba(self, ds):
         assert hasattr(
             self, "model_target_variable_name_"
         ), "Target variable name has not been set"
@@ -2382,7 +2372,17 @@ def _transform_prepare_deployment_model_for_predict_proba(self, ds):
         if self.model_target_variable_name_ not in list(ds.main_table.column_ids):
             model_dictionary.remove_variable(self.model_target_variable_name_)
 
-        return model_copy
+        if ds.is_in_memory:
+            output_columns_dtype = {}
+            if self.is_multitable_model_:
+                output_columns_dtype.update(self._original_key_dtypes)
+            for variable in model_dictionary.variables:
+                if variable.used and variable.name not in model_dictionary.key:
+                    output_columns_dtype[variable.name] = np.float64
+        else:
+            output_columns_dtype = None
+
+        return model_copy, output_columns_dtype
 
 
 # Note: scikit-learn **requires** inherit first the mixins and then other classes
@@ -2534,6 +2534,8 @@ def __init__(
         )
         self._khiops_model_prefix = "SNB_"
         self._predicted_target_meta_data_tag = "Mean"
+        self._predicted_target_name_prefix = "M"
+        self._original_target_dtype = np.float64
 
     def fit(self, X, y=None, **kwargs):
         """Fits a Selective Naive Bayes regressor according to X, y
@@ -3090,7 +3092,7 @@ def transform(self, X):
             X_transformed = super()._transform(
                 ds,
                 computation_dir,
-                self._transform_prepare_deployment_model,
+                self._transform_prepare_deployment_for_transform,
                 True,
                 "transform.txt",
             )
@@ -3102,7 +3104,7 @@ def transform(self, X):
             return X_transformed.to_numpy(copy=False)
         return X_transformed
 
-    def _transform_prepare_deployment_model(self, ds):
+    def _transform_prepare_deployment_for_transform(self, ds):
         assert hasattr(
             self, "model_target_variable_name_"
         ), "Target variable name has not been set"
@@ -3115,7 +3117,11 @@ def _transform_prepare_deployment_model(self, ds):
         if self.model_target_variable_name_ not in list(ds.main_table.column_ids):
             model_dictionary.remove_variable(self.model_target_variable_name_)
 
-        return model_copy
+        # TODO
+        #  Replace the second return value (the output columns' dtypes) with a proper
+        #  value instead of `None`. In the current state, it will use pandas type
+        #  auto-detection to load the internal table into memory.
+        return model_copy, None
 
     def fit_transform(self, X, y=None, **kwargs):
         """Fit and transforms its inputs