From 86efa339b3f72e7ddfdec5daf96edc9126367bd3 Mon Sep 17 00:00:00 2001 From: elvis-xiao Date: Sun, 24 Dec 2023 18:48:41 +0800 Subject: [PATCH] format imputer cols_transform_value Signed-off-by: elvis-xiao --- python/federatedml/feature/imputer.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/federatedml/feature/imputer.py b/python/federatedml/feature/imputer.py index 2da408a1ee..8551d26d34 100644 --- a/python/federatedml/feature/imputer.py +++ b/python/federatedml/feature/imputer.py @@ -215,7 +215,16 @@ def __get_cols_transform_value(self, data, replace_method, replace_value=None): LOGGER.debug(f"replace value for feature {feature} is: {transform_value}") else: raise ValueError("Unknown replace method:{}".format(replace_method)) - cols_transform_value[feature] = transform_value + if (data.schema.get("meta",{}).get("exclusive_data_type",{}).get(feature) is not None): + data_type = data.schema["meta"]["exclusive_data_type"][feature] + if ("str" == data_type): + cols_transform_value[feature] = str(transform_value) + elif ("int" == data_type): + cols_transform_value[feature] = int(transform_value) + else: + cols_transform_value[feature] = float(transform_value) + else: + cols_transform_value[feature] = transform_value LOGGER.debug(f"cols_transform value is: {cols_transform_value}") cols_transform_value = [cols_transform_value[key] for key in header]