blob: 5cc5229a8bed209d16580053091590a1a0b6433b (
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
|
#include "tensorflow/python/lib/io/py_record_reader.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/platform/port.h"
#include "tensorflow/core/public/env.h"
namespace tensorflow {
class RandomAccessFile;
namespace io {
PyRecordReader::PyRecordReader() {}
PyRecordReader* PyRecordReader::New(const string& filename,
uint64 start_offset) {
RandomAccessFile* file;
Status s = Env::Default()->NewRandomAccessFile(filename, &file);
if (!s.ok()) {
return nullptr;
}
PyRecordReader* reader = new PyRecordReader;
reader->offset_ = start_offset;
reader->file_ = file;
reader->reader_ = new RecordReader(reader->file_);
return reader;
}
PyRecordReader::~PyRecordReader() {
delete reader_;
delete file_;
}
bool PyRecordReader::GetNext() {
if (reader_ == nullptr) return false;
Status s = reader_->ReadRecord(&offset_, &record_);
return s.ok();
}
void PyRecordReader::Close() {
delete reader_;
delete file_;
file_ = nullptr;
reader_ = nullptr;
}
} // namespace io
} // namespace tensorflow
|
