From d95a2e9ccfd5922bfe3a14cabb7e8c5dc037e6bb Mon Sep 17 00:00:00 2001
From: ms_yan <6576637+ms_yan@user.noreply.gitee.com>
Date: Thu, 7 May 2020 23:00:53 +0800
Subject: [PATCH] repair num of input columns and output columns check

---
 mindspore/ccsrc/dataset/engine/datasetops/rename_op.cc | 7 +------
 mindspore/dataset/engine/validators.py                 | 8 ++++++++
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/mindspore/ccsrc/dataset/engine/datasetops/rename_op.cc b/mindspore/ccsrc/dataset/engine/datasetops/rename_op.cc
index 5f354abb04..92639b8809 100644
--- a/mindspore/ccsrc/dataset/engine/datasetops/rename_op.cc
+++ b/mindspore/ccsrc/dataset/engine/datasetops/rename_op.cc
@@ -51,12 +51,7 @@ Status RenameOp::Builder::Build(std::shared_ptr<RenameOp> *ptr) {
 //  constructor
 RenameOp::RenameOp(const std::vector<std::string> &in_col_names, const std::vector<std::string> &out_col_names,
                    int32_t op_connector_size)
-    : PipelineOp(op_connector_size), in_columns_(in_col_names), out_columns_(out_col_names) {
-  // check input & output sizes
-  if (in_columns_.size() != out_columns_.size()) {
-    MS_LOG(ERROR) << "Rename operator number of in columns != number of out columns.";
-  }
-}
+    : PipelineOp(op_connector_size), in_columns_(in_col_names), out_columns_(out_col_names) {}
 
 // destructor
 RenameOp::~RenameOp() {}
diff --git a/mindspore/dataset/engine/validators.py b/mindspore/dataset/engine/validators.py
index 4f1bb2c2d7..b24d647fd5 100644
--- a/mindspore/dataset/engine/validators.py
+++ b/mindspore/dataset/engine/validators.py
@@ -842,6 +842,14 @@ def check_rename(method):
                 raise ValueError("{} is not provided.".format(param_name))
             check_columns(param, param_name)
 
+        input_size, output_size = 1, 1
+        if isinstance(param_dict.get(req_param_columns[0]), list):
+            input_size = len(param_dict.get(req_param_columns[0]))
+        if isinstance(param_dict.get(req_param_columns[1]), list):
+            output_size = len(param_dict.get(req_param_columns[1]))
+        if input_size != output_size:
+            raise ValueError("Number of column in input_columns and output_columns is not equal.")
+
         return method(*args, **kwargs)
 
     return new_method