!2473 Compute column name id map during tree prepare phase

Merge pull request !2473 from lixiachen/compute_column_map

mindspore/ccsrc/dataset/engine/dataset_iterator.cc

@@ -27,7 +27,7 @@
 namespace mindspore {
 namespace dataset {
 // Constructor of the IteratorBase
-IteratorBase::IteratorBase() : curr_buffer_(nullptr), eof_handled_(false), first_row_(true) {}
+IteratorBase::IteratorBase() : curr_buffer_(nullptr), eof_handled_(false) {}
 
 IteratorBase::~IteratorBase() = default;
 
@@ -51,13 +51,10 @@ Status IteratorBase::GetNextAsMap(TensorMap *out_map) {
   // The column name mapping comes from the source operator that is producing the data into the iterator.
   // To avoid having to fetch this for every time, we'll take a local copy of the column name id mapping
   // and save in the iterator.  We only have to do this once.  All subsequent iterations use the same mapping.
-  // Note: This can only be done after the first row has been produced, as this guarantees the the child has
-  // it's column mapping set up.
-  if (first_row_) {
+  if (col_name_id_map_.empty()) {
     // Determine the column name map by calling the derived class method to retrieve the column
     // name map
     col_name_id_map_ = this->GetColumnNameMap();
-    first_row_ = false;
   }
 
   // Populate the out map from the row and return it

mindspore/ccsrc/dataset/engine/dataset_iterator.h

@@ -72,7 +72,6 @@ class IteratorBase {
  protected:
   std::unique_ptr<DataBuffer> curr_buffer_;  // holds the current buffer
   bool eof_handled_;                         // T/F if this op got an eof
-  bool first_row_;                           // internal tracking for first row case
   std::unordered_map<std::string, int32_t> col_name_id_map_;
 };
 

mindspore/ccsrc/dataset/engine/datasetops/barrier_op.cc

@@ -144,9 +144,6 @@ Status BarrierOp::prepare(TensorQTable *const table) {
 
   table->push_back(std::move(new_row));
 
-  // Assign the column name id map
-  RETURN_IF_NOT_OK(DatasetOp::AssignColMapFromChild());
-
   // the update code below shouldn't do anything bad if the column name already exists.
   return Status::OK();
 }

mindspore/ccsrc/dataset/engine/datasetops/batch_op.cc

@@ -76,7 +76,6 @@ Status BatchOp::operator()() {
   std::unique_ptr<TensorQTable> table = std::make_unique<TensorQTable>();
   child_iterator_ = std::make_unique<ChildIterator>(this, 0, 0);
   RETURN_IF_NOT_OK(child_iterator_->FetchNextTensorRow(&new_row));
-  RETURN_IF_NOT_OK(DatasetOp::AssignColMapFromChild());  // must come after the first fetch above
   int32_t cur_batch_size = 0;
   RETURN_IF_NOT_OK(GetBatchSize(&cur_batch_size, CBatchInfo(0, 0, 0)));
   while (child_iterator_->eof_handled() == false) {

mindspore/ccsrc/dataset/engine/datasetops/bucket_batch_by_length_op.cc

@@ -115,7 +115,6 @@ Status BucketBatchByLengthOp::operator()() {
   TensorRow current_row;
   child_iterator_ = std::make_unique<ChildIterator>(this, 0, 0);
   RETURN_IF_NOT_OK(child_iterator_->FetchNextTensorRow(&current_row));
-  RETURN_IF_NOT_OK(AssignColMapFromChild());
   while (!child_iterator_->eof_handled()) {
     while (!current_row.empty()) {
       int32_t element_length;

mindspore/ccsrc/dataset/engine/datasetops/build_vocab_op.cc

@@ -86,7 +86,6 @@ Status BuildVocabOp::operator()() {
   child_iterator_ = std::make_unique<ChildIterator>(this, 0, 0);
   TensorRow new_row;
   RETURN_IF_NOT_OK(child_iterator_->FetchNextTensorRow(&new_row));
-  RETURN_IF_NOT_OK(AssignColMapFromChild());
   if (!col_names_.empty()) {
     col_ids_.reserve(col_names_.size());
     for (std::string col : col_names_) {

mindspore/ccsrc/dataset/engine/datasetops/concat_op.cc

@@ -66,12 +66,6 @@ Status ConcatOp::operator()() {
   std::unique_ptr<DataBuffer> buf;
   RETURN_IF_NOT_OK(child_[0]->GetNextBuffer(&buf));
 
-  // Obtain columns_name_id_map from child_[0]
-  column_name_id_map_ = child_[0]->column_name_id_map();
-  if (column_name_id_map_.empty()) {
-    RETURN_STATUS_UNEXPECTED("Child column name map cannot be empty!");
-  }
-
   int eof_count = 0;
   while (eof_count != children_num_) {
     for (int i = 0; i < children_num_; i++) {
@@ -115,17 +109,13 @@ Status ConcatOp::Verify(int32_t id, const std::unique_ptr<DataBuffer> &buf) {
   buf->GetRow(0, &new_row);
 
   if (id == 0) {
-    // Obtain the column name, data type and data rank in child[0]
-    column_name_id_ = child_[id]->column_name_id_map();
+    // Obtain the data type and data rank in child[0]
     for (auto item : new_row) {
       data_type_.push_back(item->type());
       data_rank_.push_back(item->Rank());
     }
   } else {
-    // Compare the column name, data type and data rank with these in child[0]
-    if (child_[id]->column_name_id_map() != column_name_id_) {
-      RETURN_STATUS_UNEXPECTED("The column name or column order is not the same with previous dataset.");
-    }
+    // Compare the data type and data rank with these in child[0]
     int32_t index = 0;
     for (auto item : new_row) {
       if ((item->type() != data_type_[index]) || item->Rank() != data_rank_[index++]) {
@@ -141,5 +131,25 @@ Status ConcatOp::PrepareNodePostAction() {
   tree_->AddToRepeatStack(shared_from_this());
   return Status::OK();
 }
+
+// We need to overwrite the super class ComputeColMap here because the number of children is more than 1.
+Status ConcatOp::ComputeColMap() {
+  if (column_name_id_map_.empty()) {
+    // Obtain columns_name_id_map from child_[0]
+    column_name_id_map_ = child_[0]->column_name_id_map();
+    if (column_name_id_map_.empty()) {
+      RETURN_STATUS_UNEXPECTED("Child column name map cannot be empty!");
+    }
+    // Verify all children have the same column name map
+    for (int32_t i = 0; i < child_.size(); ++i) {
+      if (child_[i]->column_name_id_map() != column_name_id_map_) {
+        RETURN_STATUS_UNEXPECTED("The column name or column order is not the same with previous dataset.");
+      }
+    }
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
+  }
+  return Status::OK();
+}
 }  // namespace dataset
 }  // namespace mindspore

mindspore/ccsrc/dataset/engine/datasetops/concat_op.h

@@ -85,6 +85,10 @@ class ConcatOp : public PipelineOp {
   // @return Name of the current Op
   std::string Name() const override { return "ConcatOp"; }
 
+  // Private function for computing the assignment of the column name map.
+  // @return - Status
+  Status ComputeColMap() override;
+
  private:
   Status Verify(int32_t id, const std::unique_ptr<DataBuffer> &buf);
 

mindspore/ccsrc/dataset/engine/datasetops/dataset_op.cc

@@ -39,8 +39,7 @@ DatasetOp::DatasetOp(int32_t op_connector_size)
       tree_(nullptr),
       state_(OpState::kDeOpIdle),
       op_ctrl_flags_(kDeOpNone),
-      out_connector_(nullptr),
-      first_fetch_(true) {
+      out_connector_(nullptr) {
   // The operator starts out with an invalid operator id.  The only way to
   // get it out of invalid state is to assign the operator to an execution tree.
 }
@@ -240,6 +239,10 @@ Status DatasetOp::PrepareNodePostAction() {
     RETURN_IF_NOT_OK(out_connector_->Register(tree_->AllTasks()));
   }
   RETURN_IF_NOT_OK(this->RegisterWorkerConnectors());
+
+  // Generate the column name map for the current op.
+  RETURN_IF_NOT_OK(this->ComputeColMap());
+
   return Status::OK();
 }
 
@@ -262,30 +265,21 @@ std::string DatasetOp::ColumnNameMapAsString() const {
   return outStr;
 }
 
-// A helper function for providing assignment of the column name map.
-// This grabs the map from child 0 and assigns it into this op.
-// Can only be used if number of children is 1.
-Status DatasetOp::AssignColMapFromChild() {
+// Computing the assignment of the column name map.
+// This just inherits the column map from its first child, can only be used if the number of children is 1.
+// Operations changing the column map must overwrite this function.
+Status DatasetOp::ComputeColMap() {
   if (child_.size() > 1) {
     RETURN_STATUS_UNEXPECTED("Assigning column name map from child only works for single-child operators.");
   }
-  // Assign the correct column name map to this op by taking it from the input child.
-  // This must be done AFTER the first fetch, but only needs to be done once by the first worker to
-  // do the first fetch.
-  if (first_fetch_) {
-    // If there was a single worker, or this is being called from a master thread in a parallel op,
-    // then the mutex is not really needed here, although it's harmless.
-    std::unique_lock<std::mutex> lock(column_name_map_mutex_);
-    // If the map has not been set up yet, then we are the first one in to set it up. The first_fetch_ (dirty read)
-    // bool allows us to avoid acquiring the lock if the map has already been set.
   if (column_name_id_map_.empty()) {
     column_name_id_map_ = child_[0]->column_name_id_map();
-      first_fetch_ = false;
     if (column_name_id_map_.empty()) {
       RETURN_STATUS_UNEXPECTED("Child column name map cannot be empty!");
     }
-    }
-    MS_LOG(DEBUG) << "Setting column map after first fetch:\n" << DatasetOp::ColumnNameMapAsString();
+    MS_LOG(DEBUG) << "Setting column map:\n" << DatasetOp::ColumnNameMapAsString();
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
   }
   return Status::OK();
 }

mindspore/ccsrc/dataset/engine/datasetops/dataset_op.h

@@ -277,11 +277,12 @@ class DatasetOp : public std::enable_shared_from_this<DatasetOp> {
   // @param parent - The parent node to remove
   void RemoveParent(DatasetOp *parent);
 
-  // A helper function for providing an assignment of the column name map.
-  // This grabs the map from child 0 and assigns it into this op.
-  // Can only be used if number of children is 1.
+  // Compute the current op's column map using its child's column map.
+  // Get called during the tree post-prepare phase in PrepareNodePostAction.
+  // This base implementation just inherits the map from child 0, and can only be used if the number of children is 1.
+  // Operations changing the column map it inherits from the child must overwrite this function.
   // @return - Status
-  Status AssignColMapFromChild();
+  virtual Status ComputeColMap();
 
   std::vector<std::shared_ptr<DatasetOp>> child_;                // Child nodes
   std::vector<DatasetOp *> parent_;                              // Parent nodes. No ownership
@@ -292,7 +293,6 @@ class DatasetOp : public std::enable_shared_from_this<DatasetOp> {
   uint32_t op_ctrl_flags_;                                       // Flags for the operator
   std::unique_ptr<DbConnector> out_connector_;                   // Output Connector
   std::unordered_map<std::string, int32_t> column_name_id_map_;  // Mapping between col index and col name
-  bool first_fetch_;                                             // For use when setting column map
   std::mutex column_name_map_mutex_;                             // For protecting shared access to the column map
 
  private:

mindspore/ccsrc/dataset/engine/datasetops/filter_op.cc

@@ -126,9 +126,6 @@ Status FilterOp::WorkerEntry(int32_t worker_id) {
       continue;
     }
 
-    // Now that the first fetch is in, use the helper function to assign the column name map to this op.
-    RETURN_IF_NOT_OK(DatasetOp::AssignColMapFromChild());
-
     RETURN_IF_NOT_OK(CheckColumns(in_buffer.get(), &in_columns_));
 
     // if the databuffer was all filtered, it is marked as kFilterEmpty.

mindspore/ccsrc/dataset/engine/datasetops/map_op.cc

@@ -156,14 +156,15 @@ Status MapOp::WorkerEntry(int32_t worker_id) {
   // initializations that happen after the first fetch.
   RETURN_IF_NOT_OK(FetchNextBuffer(&in_buffer, worker_id));
 
-  // Initialize details related to column selections and column map by calling WorkerEntryInit.
-  // WorkerEntryInit contains thread-safe lock to ensure that this init work is only performed once
-  // by the first worker to enter the codepath. All other threads will share the const info that
-  // gets set up here going forward.
+  // Sanity check the databuffer.
   // Special case: if there's more threads than buffers, some threads simply get the final control
-  // messages (eoe/eof), and so they will not perform the init work.
+  // messages (eoe/eof), and so they will not perform the check.
   if (!in_buffer->eoe() && !in_buffer->eof()) {
-    RETURN_IF_NOT_OK(WorkerEntryInit(in_buffer.get()));
+    int32_t num_rows = in_buffer->NumRows();
+    int32_t num_cols = in_buffer->NumCols();
+    if (num_rows == 0 || num_cols == 0) {
+      RETURN_STATUS_UNEXPECTED("MapOp is getting an empty DataBuffer.");
+    }
   }
 
   // Now that init work is done, drop into the main fetching loop.
@@ -258,26 +259,37 @@ Status MapOp::WorkerCompute(DataBuffer *in_buffer, TensorQTable *new_tensor_tabl
   return Status::OK();
 }
 
-// initialize some internal data structure used by WorkerEntry()
-Status MapOp::WorkerEntryInit(const DataBuffer *in_buf) {
-  int32_t num_rows = in_buf->NumRows();
-  int32_t num_cols = in_buf->NumCols();
-  if (num_rows == 0 || num_cols == 0) {
-    RETURN_STATUS_UNEXPECTED("MapOp is getting an empty DataBuffer.");
-  }
-
-  // We can't use AssignColMapFromChild() here since we need to modify the column map. We need to be threadsafe
-  // though for saving the final map in the op, so use the lock here.
-  if (first_fetch_) {
-    std::unique_lock<std::mutex> lock(column_name_map_mutex_);
-    // If the map has not been set up yet in the base class, then we are the first one in to set it up
-    // (and we are under protection of the mutex lock)
+Status MapOp::ComputeColMap() {
+  // If the map has not been set up yet in the base class, then set it up
   if (column_name_id_map_.empty()) {
     std::unordered_map<std::string, int32_t> current_name_id_map = child_[0]->column_name_id_map();
+    // Initialize private variables
+    RETURN_IF_NOT_OK(InitPrivateVariable(&current_name_id_map));
+    // Create the final column name to index mapping in the base class field
+    CreateFinalColMap(&current_name_id_map);
+    MS_LOG(DEBUG) << "Column name map for map op set: " << this->ColumnNameMapAsString();
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
+  }
+  return Status::OK();
+}
 
+// Validating if each of the input_columns exists in the DataBuffer.
+Status MapOp::ValidateInColumns(const std::unordered_map<std::string, int32_t> &col_name_id_map) {
+  for (const auto &inCol : in_columns_) {
+    bool found = col_name_id_map.find(inCol) != col_name_id_map.end() ? true : false;
+    if (!found) {
+      std::string err_msg = "input column name: " + inCol + " doesn't exist in the dataset columns.";
+      RETURN_STATUS_UNEXPECTED(err_msg);
+    }
+  }
+  return Status::OK();
+}
+
+Status MapOp::InitPrivateVariable(std::unordered_map<std::string, int32_t> *col_name_id_map) {
   // If input_columns is empty(), The col at index-0 will be picked.
   if (in_columns_.empty()) {
-        for (const auto &pair : current_name_id_map) {
+    for (const auto &pair : *col_name_id_map) {
       if (pair.second == 0) {
         MS_LOG(INFO) << "Input columns empty for map op, will apply to the first column in the current table.";
         in_columns_.push_back(pair.first);
@@ -294,38 +306,18 @@ Status MapOp::WorkerEntryInit(const DataBuffer *in_buf) {
 
   // Before we continue, issue a sanity check to make sure the input columns from user and the incoming
   // columns from child are correct
-      RETURN_IF_NOT_OK(this->ValidateInColumns(current_name_id_map));
+  RETURN_IF_NOT_OK(this->ValidateInColumns(*col_name_id_map));
 
   // initialize keep_input_columns, true means to keep the column.
-      keep_input_columns_.resize(num_cols, true);
+  keep_input_columns_.resize(col_name_id_map->size(), true);
   for (const auto &col_name : in_columns_) {
-        int32_t missed = current_name_id_map[col_name];
+    int32_t missed = (*col_name_id_map)[col_name];
     keep_input_columns_[missed] = false;
   }
 
   // initialize to_process_indices.
   for (const auto &col_name : in_columns_) {
-        to_process_indices_.push_back(current_name_id_map[col_name]);
-      }
-
-      // Create the final column name to index mapping in the base class field
-      CreateFinalColMap(&current_name_id_map);
-      first_fetch_ = false;
-    }
-  }  // mutex lock will release here
-
-  MS_LOG(DEBUG) << "Column name map for map op set: " << this->ColumnNameMapAsString();
-  return Status::OK();
-}
-
-// Validating if each of the input_columns exists in the DataBuffer.
-Status MapOp::ValidateInColumns(const std::unordered_map<std::string, int32_t> &col_name_id_map) {
-  for (const auto &inCol : in_columns_) {
-    bool found = col_name_id_map.find(inCol) != col_name_id_map.end() ? true : false;
-    if (!found) {
-      std::string err_msg = "input column name: " + inCol + " doesn't exist in the dataset columns.";
-      RETURN_STATUS_UNEXPECTED(err_msg);
-    }
+    to_process_indices_.push_back((*col_name_id_map)[col_name]);
   }
   return Status::OK();
 }

mindspore/ccsrc/dataset/engine/datasetops/map_op.h

@@ -258,15 +258,18 @@ class MapOp : public ParallelOp {
   // @param col_name_id_map The column name to index mapping obtained from child operator
   void CreateFinalColMap(std::unordered_map<std::string, int32_t> *col_name_id_map);
 
-  // Private function that initialize some internal data structure used by WorkerEntry()
-  // @param in_buf A raw pointer to the DataBuffer. A raw pointer is fine because this function does not manage memory
-  //     and is not shared with other threads.
-  Status WorkerEntryInit(const DataBuffer *in_buf);
-
   // Validating if each of the input_columns exists in the DataBuffer.
   // @param - the column map to check
   // @return - status return code
   Status ValidateInColumns(const std::unordered_map<std::string, int32_t> &col_name_id_map);
+
+  // Private function for computing the assignment of the column name map.
+  // @return - Status
+  Status ComputeColMap() override;
+
+  // Private function for initializing private variables such as in_columns_, out_columns_.
+  // @return - Status
+  Status InitPrivateVariable(std::unordered_map<std::string, int32_t> *col_name_id_map);
 };
 }  // namespace dataset
 }  // namespace mindspore

mindspore/ccsrc/dataset/engine/datasetops/project_op.cc

@@ -74,24 +74,6 @@ void ProjectOp::Print(std::ostream &out, bool show_all) const {
 Status ProjectOp::GetNextBuffer(std::unique_ptr<DataBuffer> *p_buffer, int32_t worker_id, bool retry_if_eoe) {
   RETURN_IF_NOT_OK(child_[0]->GetNextBuffer(p_buffer, worker_id, retry_if_eoe));
   if (!((*p_buffer)->eoe()) && !((*p_buffer)->eof())) {
-    // Only for the first buffer fetched, get the column map of the incoming data and save it
-    // into our own column name map after making the appropriate mods
-    // We cannot use the super class AssignColMapFromChild here because we're making a modification of the
-    // map from the child map.
-    if (first_fetch_) {
-      std::unordered_map<std::string, int32_t> child_column_name_mapping = child_[0]->column_name_id_map();
-      for (size_t i = 0; i < columns_to_project_.size(); i++) {
-        std::string &current_column = columns_to_project_[i];
-        if (child_column_name_mapping.find(current_column) == child_column_name_mapping.end()) {
-          std::string err_msg = "ProjectOp: column " + current_column + " does not exist in child operator.";
-          RETURN_STATUS_UNEXPECTED(err_msg);
-        }
-        // Setup the new column name mapping for ourself (base class field)
-        column_name_id_map_[current_column] = i;
-        projected_column_indices_.push_back(child_column_name_mapping[current_column]);
-      }
-      first_fetch_ = false;  // we only need to do this path once
-    }
     RETURN_IF_NOT_OK(Project(p_buffer));
   }
   return Status::OK();
@@ -151,5 +133,27 @@ Status ProjectOp::Accept(NodePass *p, bool *modified) {
   // Downcast shared pointer then call visitor
   return p->RunOnNode(std::static_pointer_cast<ProjectOp>(shared_from_this()), modified);
 }
+
+// Compute the column map and save it into our own column name map
+// We cannot use the super class ComputeColMap here because we're making a modification of the
+// map from the child map.
+Status ProjectOp::ComputeColMap() {
+  if (column_name_id_map_.empty()) {
+    std::unordered_map<std::string, int32_t> child_column_name_mapping = child_[0]->column_name_id_map();
+    for (size_t i = 0; i < columns_to_project_.size(); i++) {
+      std::string &current_column = columns_to_project_[i];
+      if (child_column_name_mapping.find(current_column) == child_column_name_mapping.end()) {
+        std::string err_msg = "ProjectOp: column " + current_column + " does not exist in child operator.";
+        RETURN_STATUS_UNEXPECTED(err_msg);
+      }
+      // Setup the new column name mapping for ourself (base class field)
+      column_name_id_map_[current_column] = i;
+      projected_column_indices_.push_back(child_column_name_mapping[current_column]);
+    }
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
+  }
+  return Status::OK();
+}
 }  // namespace dataset
 }  // namespace mindspore

mindspore/ccsrc/dataset/engine/datasetops/project_op.h

@@ -116,6 +116,10 @@ class ProjectOp : public PipelineOp {
   std::vector<int32_t> projected_column_indices_;
 
   Status Project(std::unique_ptr<DataBuffer> *data_buffer);
+
+  // Computing the assignment of the column name map.
+  // @return - Status
+  Status ComputeColMap() override;
 };
 }  // namespace dataset
 }  // namespace mindspore

mindspore/ccsrc/dataset/engine/datasetops/rename_op.cc

@@ -69,12 +69,6 @@ Status RenameOp::operator()() {
     RETURN_STATUS_UNEXPECTED(err_msg);
   }
 
-  // First, populate the column map from the input child.
-  // This will not be the final map for output from this op.
-  RETURN_IF_NOT_OK(DatasetOp::AssignColMapFromChild());
-  // core rename functionality only needs to happen once, to identify the new column names/indexes
-  RETURN_IF_NOT_OK(RenameColumns());
-
   while (curr_buffer->eof() == false) {
     while (curr_buffer->eoe() == false) {
       // push the renamed input buffer
@@ -95,8 +89,12 @@ Status RenameOp::operator()() {
   return Status::OK();
 }
 
-// renames the columns
-Status RenameOp::RenameColumns() {
+// Rename core functionality to compute the new column name id map.
+// We need to overwrite the super class ComputeColMap here because we're making a modification of the
+// map from the child map.
+Status RenameOp::ComputeColMap() {
+  if (column_name_id_map_.empty()) {
+    column_name_id_map_ = child_[0]->column_name_id_map();
     // iterate over my index in input vector, find the corresponding position
     std::unordered_map<std::string, int32_t> new_col_name_id_map = {};
     // parameter for input check
@@ -134,6 +132,9 @@ Status RenameOp::RenameColumns() {
 
     // Now, overwrite our column map with the new renamed columns/id's
     column_name_id_map_ = new_col_name_id_map;
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
+  }
   return Status::OK();
 }
 

mindspore/ccsrc/dataset/engine/datasetops/rename_op.h

@@ -122,7 +122,9 @@ class RenameOp : public PipelineOp {
 
  protected:
   // Rename core functionality
-  Status RenameColumns();
+  // Computing the assignment of the new column name map.
+  // @return - Status
+  Status ComputeColMap() override;
 
   // Variable to store the input column names
   std::vector<std::string> in_columns_;

mindspore/ccsrc/dataset/engine/datasetops/repeat_op.cc

@@ -123,8 +123,6 @@ Status RepeatOp::GetNextBuffer(std::unique_ptr<DataBuffer> *p_buffer, int32_t wo
   if (buf->eof()) {
     RETURN_IF_NOT_OK(EofReceived(worker_id));
   }
-  // Update the column name map if needed
-  RETURN_IF_NOT_OK(DatasetOp::AssignColMapFromChild());
   *p_buffer = std::move(buf);
   return Status::OK();
 }

mindspore/ccsrc/dataset/engine/datasetops/shuffle_op.cc

@@ -266,9 +266,6 @@ Status ShuffleOp::InitShuffleBuffer() {
     RETURN_STATUS_UNEXPECTED("Unable to fetch a single row for shuffle buffer.");
   }
 
-  // Now that a first fetch is done, assign the column map for this operator
-  RETURN_IF_NOT_OK(DatasetOp::AssignColMapFromChild());
-
   // Now fill the rest of the shuffle buffer until we are unable to get the next row or we reached
   // the desired shuffle buffer size.
   while (!new_row.empty() && shuffle_buffer_->size() < static_cast<size_t>(shuffle_size_ - 1)) {

mindspore/ccsrc/dataset/engine/datasetops/skip_op.cc

@@ -86,9 +86,6 @@ Status SkipOp::operator()() {
   std::unique_ptr<DataBuffer> curr_buffer;
   RETURN_IF_NOT_OK(GetNextInput(&curr_buffer));
 
-  // After the first buffer fetch above we can do the one-time assign of the column name map
-  RETURN_IF_NOT_OK(DatasetOp::AssignColMapFromChild());
-
   while (curr_buffer->eof() == false) {
     // Reset count
     skip_count_ = 0;

mindspore/ccsrc/dataset/engine/datasetops/source/celeba_op.cc

@@ -79,11 +79,6 @@ CelebAOp::CelebAOp(int32_t num_workers, int32_t rows_per_buffer, const std::stri
       sampler_(std::move(sampler)),
       num_rows_in_attr_file_(0),
       dataset_type_(dataset_type) {
-  // Set the column name map (base class field)
-  for (int32_t index = 0; index < data_schema_->NumColumns(); index++) {
-    column_name_id_map_[data_schema_->column(index).name()] = index;
-  }
-
   attr_info_queue_ = std::make_unique<Queue<std::vector<std::string>>>(queue_size);
   io_block_queues_.Init(num_workers_, queue_size);
 }
@@ -413,5 +408,17 @@ Status CelebAOp::Reset() {
   wp_.Set();  // wake up master thread after reset is done
   return Status::OK();
 }
+
+Status CelebAOp::ComputeColMap() {
+  // Set the column name map (base class field)
+  if (column_name_id_map_.empty()) {
+    for (int32_t index = 0; index < data_schema_->NumColumns(); index++) {
+      column_name_id_map_[data_schema_->column(index).name()] = index;
+    }
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
+  }
+  return Status::OK();
+}
 }  // namespace dataset
 }  // namespace mindspore

mindspore/ccsrc/dataset/engine/datasetops/source/celeba_op.h

@@ -212,6 +212,10 @@ class CelebAOp : public ParallelOp, RandomAccessOp {
   // @return Status - The error code return
   Status Reset() override;
 
+  // Private function for computing the assignment of the column name map.
+  // @return - Status
+  Status ComputeColMap() override;
+
   int32_t rows_per_buffer_;
   std::string folder_path_;  // directory of celeba folder
   bool decode_;

mindspore/ccsrc/dataset/engine/datasetops/source/cifar_op.cc

@@ -87,10 +87,6 @@ CifarOp::CifarOp(CifarType type, int32_t num_works, int32_t rows_per_buf, const
       sampler_(std::move(sampler)),
       row_cnt_(0),
       buf_cnt_(0) {
-  // set the column name map (base class field)
-  for (uint32_t i = 0; i < data_schema_->NumColumns(); ++i) {
-    column_name_id_map_[data_schema_->column(i).name()] = i;
-  }
   constexpr uint64_t kUtilQueueSize = 512;
   cifar_raw_data_block_ = std::make_unique<Queue<std::vector<unsigned char>>>(kUtilQueueSize);
   io_block_queues_.Init(num_workers_, queue_size);
@@ -454,5 +450,17 @@ Status CifarOp::CountTotalRows(const std::string &dir, bool isCIFAR10, int64_t *
     return Status::OK();
   }
 }
+
+Status CifarOp::ComputeColMap() {
+  // set the column name map (base class field)
+  if (column_name_id_map_.empty()) {
+    for (uint32_t i = 0; i < data_schema_->NumColumns(); ++i) {
+      column_name_id_map_[data_schema_->column(i).name()] = i;
+    }
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
+  }
+  return Status::OK();
+}
 }  // namespace dataset
 }  // namespace mindspore

mindspore/ccsrc/dataset/engine/datasetops/source/cifar_op.h

@@ -208,6 +208,10 @@ class CifarOp : public ParallelOp, public RandomAccessOp {
   // @return Status - The error code return
   Status GetClassIds(std::map<int32_t, std::vector<int64_t>> *cls_ids) const override;
 
+  // Private function for computing the assignment of the column name map.
+  // @return - Status
+  Status ComputeColMap() override;
+
   CifarType cifar_type_;
   int32_t rows_per_buffer_;
   std::string folder_path_;

mindspore/ccsrc/dataset/engine/datasetops/source/clue_op.cc

@@ -112,13 +112,6 @@ Status ClueOp::Init() {
   int32_t safe_queue_size = static_cast<int32_t>(std::ceil(clue_files_list_.size() / num_workers_) + 1);
   io_block_queues_.Init(num_workers_, safe_queue_size);
 
-  // Set the column name mapping (base class field)
-  int count = 0;
-  for (auto &p : cols_to_keyword_) {
-    column_name_id_map_[p.first] = count;
-    count++;
-  }
-
   RETURN_IF_NOT_OK(ParallelOp::CreateWorkerConnector(worker_connector_size_));
   jagged_buffer_connector_ = std::make_unique<JaggedConnector>(num_workers_, 1, worker_connector_size_);
 
@@ -549,5 +542,19 @@ Status ClueOp::CountAllFileRows(const std::vector<std::string> &files, int64_t *
   }
   return Status::OK();
 }
+
+Status ClueOp::ComputeColMap() {
+  // Set the column name mapping (base class field)
+  if (column_name_id_map_.empty()) {
+    int count = 0;
+    for (auto &p : cols_to_keyword_) {
+      column_name_id_map_[p.first] = count;
+      count++;
+    }
+  } else {
+    MS_LOG(WARNING) << "Column name map is already set!";
+  }
+  return Status::OK();
+}
 }  // namespace dataset
 }  // namespace mindspore