From e6c25859d4de6fa40b6c8f983a590a89d9f637b6 Mon Sep 17 00:00:00 2001
From: ms_yan <6576637+ms_yan@user.noreply.gitee.com>
Date: Tue, 19 May 2020 20:27:23 +0800
Subject: [PATCH] check wether input column exist in TFRecord files

---
 .../ccsrc/dataset/engine/datasetops/source/tf_reader_op.cc     | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mindspore/ccsrc/dataset/engine/datasetops/source/tf_reader_op.cc b/mindspore/ccsrc/dataset/engine/datasetops/source/tf_reader_op.cc
index 556aa53c73..12b10e6997 100644
--- a/mindspore/ccsrc/dataset/engine/datasetops/source/tf_reader_op.cc
+++ b/mindspore/ccsrc/dataset/engine/datasetops/source/tf_reader_op.cc
@@ -918,6 +918,9 @@ Status TFReaderOp::CreateSchema(const std::string tf_file, const std::vector<std
                          [](const auto &it) -> std::string { return it.first; });
   for (const auto &curr_col_name : columns) {
     auto it = feature_map.find(curr_col_name);
+    if (it == feature_map.end()) {
+      RETURN_STATUS_UNEXPECTED("Failed to find column " + curr_col_name);
+    }
     std::string column_name = it->first;
 
     std::string column_type;