aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/framework/reader_interface.h
blob: b307c37f015f1deb5877ae8aa83babe043aaaaff (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
#ifndef TENSORFLOW_FRAMEWORK_READER_INTERFACE_H_
#define TENSORFLOW_FRAMEWORK_READER_INTERFACE_H_

#include <memory>
#include <string>
#include <vector>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/platform/port.h"
#include "tensorflow/core/public/status.h"
#include "tensorflow/core/public/tensor.h"

namespace tensorflow {

class QueueInterface;
class ReaderInterface;

// Readers are the mechanism for reading records from files in
// TensorFlow graphs.  Each supported file format has a corresponding
// ReaderInterface descendant and a corresponding Op & OpKernel
// (implemented using ReaderOpKernel from reader_op_kernel.h).
//
// To use a Reader, you first encode "work" (some string, typically a
// filename) in the Reader's "work queue".  It then processes the
// "work" (reading records from the file), to produce key/value
// strings.  The methods of this class are called by ReaderFoo ops,
// so see ../ops/io_ops.cc for detailed descriptions.
//
// All descendants of this class must be thread-safe.
//
// See the design document here:
// https://docs.google.com/document/d/1UAgZOoeehYr20TdzW2CoZ30V-aqQphU4SwKXsW7eJv4/edit#

// TODO(josh11b): Switch this to Async.
class ReaderInterface : public ResourceBase {
 public:
  // Read a single record into *key / *value.  May get more work from
  // *queue if the current work is complete.  Sets the status on
  // *context with an OutOfRange Status if the the current work is
  // complete and the queue is done (closed and empty).
  // This method may block.
  virtual void Read(QueueInterface* queue, string* key, string* value,
                    OpKernelContext* context) = 0;

  // Restore this reader to its newly-constructed state.
  virtual Status Reset() = 0;

  // Accessors
  virtual int64 NumRecordsProduced() = 0;
  virtual int64 NumWorkUnitsCompleted() = 0;

  // -- Serialization/Restoration support --
  // Not all readers will support saving and restoring state.
  virtual Status SerializeState(string* state) = 0;
  // Note: Must Reset on error.
  virtual Status RestoreState(const string& state) = 0;

  string DebugString() override { return "a reader"; }

 protected:
  virtual ~ReaderInterface() {}
};

}  // namespace tensorflow

#endif  // TENSORFLOW_FRAMEWORK_READER_INTERFACE_H_