1 /** 2 * Copyright 2020-2021 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_DATASETOPS_TAKE_OP_H_ 17 #define MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_DATASETOPS_TAKE_OP_H_ 18 19 #include <memory> 20 #include <string> 21 #include <vector> 22 #include "minddata/dataset/engine/datasetops/pipeline_op.h" 23 #include "minddata/dataset/engine/dataset_iterator.h" 24 25 namespace mindspore { 26 namespace dataset { 27 class TakeOp : public PipelineOp { 28 public: 29 // Constructor of the TakeOp. 30 // @note The builder class should be used to call it 31 // @param count - The number of takes to do 32 explicit TakeOp(int32_t count); 33 34 // Destructor 35 ~TakeOp() = default; 36 37 // A print method typically used for debugging 38 // \param[in] out - The output stream to write output to 39 // \param[in] show_all - A bool to control if you want to show all info or just a summary 40 void Print(std::ostream &out, bool show_all) const override; 41 42 // << Stream output operator overload 43 // @notes This allows you to write the debug print info using stream operators 44 // \param[in] out - reference to the output stream being overloaded 45 // \param[in] ro - reference to the TakeOp to display 46 // @return - the output stream must be returned 47 friend std::ostream &operator<<(std::ostream &out, const TakeOp &ro) { 48 ro.Print(out, false); 49 return out; 50 } 51 52 // All dataset ops operate by launching a thread (see ExecutionTree). This class functor will 53 // provide the master loop that drives the logic for performing the work 54 // @return Status The status code returned 55 Status operator()() override; 56 57 // Op name getter 58 // @return Name of the current Op Name()59 std::string Name() const override { return kTakeOp; } 60 61 Status GetNextRow(TensorRow *row, int32_t worker_id, bool retry_if_eoe) override; 62 int32_t NumConsumers() const override; 63 int32_t NumProducers() const override; 64 65 private: 66 int32_t max_takes_; // The number of takes that the user requested 67 int32_t take_count_; // A counter for the current number of executed takes 68 69 std::unique_ptr<ChildIterator> child_iterator_; // An iterator for fetching. 70 }; 71 } // namespace dataset 72 } // namespace mindspore 73 74 #endif // MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_DATASETOPS_TAKE_OP_H_ 75