Add seek to FileIO via RandomAccessInputStream and BufferedInputStream,

by introducing a Seek() API to those implementations.  Seek is different
than SkipNBytes in that you can Seek backwards, whereas on a stream,
you can only move forwards.  While Seek() on RandomAccessInputStream
sort of makes sense, Seek() on BufferedInputStream makes a little less
sense.  I avoided adding Seek() to the InputStreamInterface because I didn't
want all implementors of InputStreamInterface to have to implement Seek(),
only those for which it made some sense.  However, I did have to add
a Reset() API to InputStreamInterface to allow reseting + skipping as the
method of seeking backwards.

Adds read(n) to FileIO.  Defaults to 'read to the end of the file' as
usual, and allows reading n bytes.

This change is to help enable FileIO to be used as a fileobj for gzipfile,
among other things.
Change: 134542845

1 files changed, 30 insertions, 0 deletions

diff --git a/tensorflow/core/lib/io/buffered_inputstream_test.cc b/tensorflow/core/lib/io/buffered_inputstream_test.cc
index 1263f0f9aa..4488637cad 100644
--- a/tensorflow/core/lib/io/buffered_inputstream_test.cc
+++ b/tensorflow/core/lib/io/buffered_inputstream_test.cc
@@ -293,6 +293,36 @@ TEST(BufferedInputStream, SkipNBytesRandomAccessFile) {
   }
 }
 
+TEST(BufferedInputStream, Seek) {
+  Env* env = Env::Default();
+  string fname = testing::TmpDir() + "/buffered_inputstream_test";
+  WriteStringToFile(env, fname, "0123456789");
+  std::unique_ptr<RandomAccessFile> file;
+  TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
+
+  for (auto buf_size : BufferSizes()) {
+    std::unique_ptr<RandomAccessInputStream> input_stream(
+        new RandomAccessInputStream(file.get()));
+    string read;
+    BufferedInputStream in(input_stream.get(), buf_size);
+
+    // Seek forward
+    TF_ASSERT_OK(in.Seek(3));
+    EXPECT_EQ(3, in.Tell());
+
+    // Read 4 bytes
+    TF_ASSERT_OK(in.ReadNBytes(4, &read));
+    EXPECT_EQ(read, "3456");
+    EXPECT_EQ(7, in.Tell());
+
+    // Seek backwards
+    TF_ASSERT_OK(in.Seek(1));
+    TF_ASSERT_OK(in.ReadNBytes(4, &read));
+    EXPECT_EQ(read, "1234");
+    EXPECT_EQ(5, in.Tell());
+  }
+}
+
 }  // anonymous namespace
 }  // namespace io
 }  // namespace tensorflow