diff options
Diffstat (limited to 'tensorflow/core/kernels/skip_dataset_op.cc')
-rw-r--r-- | tensorflow/core/kernels/skip_dataset_op.cc | 63 |
1 files changed, 4 insertions, 59 deletions
diff --git a/tensorflow/core/kernels/skip_dataset_op.cc b/tensorflow/core/kernels/skip_dataset_op.cc index 7ee945dd4c..52a6116a7c 100644 --- a/tensorflow/core/kernels/skip_dataset_op.cc +++ b/tensorflow/core/kernels/skip_dataset_op.cc @@ -35,14 +35,14 @@ class SkipDatasetOp : public UnaryDatasetOpKernel { int64 count; OP_REQUIRES_OK(ctx, ParseScalarArgument<int64>(ctx, "count", &count)); - *output = new Dataset(ctx, count, input); + *output = new Dataset(count, input); } private: - class Dataset : public GraphDatasetBase { + class Dataset : public DatasetBase { public: - Dataset(OpKernelContext* ctx, int64 count, const DatasetBase* input) - : GraphDatasetBase(ctx), count_(count), input_(input) { + Dataset(int64 count, const DatasetBase* input) + : count_(count), input_(input) { input_->Ref(); } @@ -71,18 +71,6 @@ class SkipDatasetOp : public UnaryDatasetOpKernel { string DebugString() override { return "SkipDatasetOp::Dataset"; } - protected: - Status AsGraphDefInternal(DatasetGraphDefBuilder* b, - Node** output) const override { - Node* input_graph_node = nullptr; - TF_RETURN_IF_ERROR(b->AddParentDataset(input_, &input_graph_node)); - Node* count = nullptr; - TF_RETURN_IF_ERROR(b->AddScalar(count_, &count)); - TF_RETURN_IF_ERROR( - b->AddDataset(this, {input_graph_node, count}, output)); - return Status::OK(); - } - private: class EmptyIterator : public DatasetIterator<Dataset> { public: @@ -94,16 +82,6 @@ class SkipDatasetOp : public UnaryDatasetOpKernel { *end_of_sequence = true; return Status::OK(); } - - protected: - Status SaveInternal(IteratorStateWriter* writer) override { - return Status::OK(); - } - - Status RestoreInternal(OpKernelContext* ctx, - IteratorStateReader* reader) override { - return Status::OK(); - } }; class FiniteIterator : public DatasetIterator<Dataset> { @@ -118,11 +96,6 @@ class SkipDatasetOp : public UnaryDatasetOpKernel { bool* end_of_sequence) override { mutex_lock l(mu_); // TODO(mrry): Make locking less conservative. - if (!input_impl_) { - *end_of_sequence = true; - return Status::OK(); - } - // Keep calling GetNext(). TODO(vrv): Figure out a way to // skip records without reading, perhaps by adding an // interface to iterator. @@ -143,34 +116,6 @@ class SkipDatasetOp : public UnaryDatasetOpKernel { // Return GetNext() on the underlying iterator. TF_RETURN_IF_ERROR(input_impl_->GetNext(ctx, out_tensors, end_of_sequence)); - if (*end_of_sequence) { - input_impl_.reset(); - } - return Status::OK(); - } - - protected: - Status SaveInternal(IteratorStateWriter* writer) override { - mutex_lock l(mu_); - TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("i"), i_)); - if (input_impl_) { - TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_)); - } else { - TF_RETURN_IF_ERROR( - writer->WriteScalar(full_name("input_impl_empty"), "")); - } - return Status::OK(); - } - - Status RestoreInternal(OpKernelContext* ctx, - IteratorStateReader* reader) override { - mutex_lock l(mu_); - TF_RETURN_IF_ERROR(reader->ReadScalar(full_name("i"), &i_)); - if (!reader->Contains(full_name("input_impl_empty"))) { - TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_)); - } else { - input_impl_.reset(); - } return Status::OK(); } |