1 | /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | #ifndef TENSORFLOW_CORE_FRAMEWORK_READER_INTERFACE_H_ |
17 | #define TENSORFLOW_CORE_FRAMEWORK_READER_INTERFACE_H_ |
18 | |
19 | #include <memory> |
20 | #include <string> |
21 | #include "tensorflow/core/framework/op_kernel.h" |
22 | #include "tensorflow/core/framework/resource_mgr.h" |
23 | #include "tensorflow/core/framework/tensor.h" |
24 | #include "tensorflow/core/lib/core/status.h" |
25 | #include "tensorflow/core/platform/types.h" |
26 | |
27 | namespace tensorflow { |
28 | |
29 | class QueueInterface; |
30 | class ReaderInterface; |
31 | |
32 | // Readers are the mechanism for reading records from files in |
33 | // TensorFlow graphs. Each supported file format has a corresponding |
34 | // ReaderInterface descendant and a corresponding Op & OpKernel |
35 | // (implemented using ReaderOpKernel from reader_op_kernel.h). |
36 | // |
37 | // To use a Reader, you first encode "work" (some string, typically a |
38 | // filename) in the Reader's "work queue". It then processes the |
39 | // "work" (reading records from the file), to produce key/value |
40 | // strings. The methods of this class are called by ReaderFoo ops, |
41 | // so see ../ops/io_ops.cc for detailed descriptions. |
42 | // |
43 | // All descendants of this class must be thread-safe. |
44 | class ReaderInterface : public ResourceBase { |
45 | public: |
46 | // Read a single record into *key / *value. May get more work from |
47 | // *queue if the current work is complete. Sets the status on |
48 | // *context with an OutOfRange Status if the current work is |
49 | // complete and the queue is done (closed and empty). |
50 | // This method may block. |
51 | virtual void Read(QueueInterface* queue, tstring* key, tstring* value, |
52 | OpKernelContext* context) = 0; |
53 | |
54 | // Read up to num_records records into keys / values. May get more work from |
55 | // *queue if the current work is complete. Sets the status on |
56 | // *context with an OutOfRange Status if the current work is |
57 | // complete and the queue is done (closed and empty). |
58 | // This method may block. |
59 | // The std::vector keys/value pointers are assumed to point to empty |
60 | // structures (that have most likely been reserve(num_records)). |
61 | // Returns how many records were actually read. |
62 | virtual int64_t ReadUpTo(const int64_t num_records, QueueInterface* queue, |
63 | std::vector<tstring>* keys, |
64 | std::vector<tstring>* value, |
65 | OpKernelContext* context) = 0; |
66 | |
67 | // Restore this reader to its newly-constructed state. |
68 | virtual Status Reset() = 0; |
69 | |
70 | // Accessors |
71 | virtual int64_t NumRecordsProduced() = 0; |
72 | virtual int64_t NumWorkUnitsCompleted() = 0; |
73 | |
74 | // -- Serialization/Restoration support -- |
75 | // Not all readers will support saving and restoring state. |
76 | virtual Status SerializeState(tstring* state) = 0; |
77 | // Note: Must Reset on error. |
78 | virtual Status RestoreState(const tstring& state) = 0; |
79 | |
80 | string DebugString() const override { return "a reader" ; } |
81 | |
82 | protected: |
83 | virtual ~ReaderInterface() {} |
84 | }; |
85 | |
86 | } // namespace tensorflow |
87 | |
88 | #endif // TENSORFLOW_CORE_FRAMEWORK_READER_INTERFACE_H_ |
89 | |