Add a ReaderReadUpTo Op

- Add a ReadUpTo method to ReaderInterface
- Add ReadUpToLocked to ReaderBase w/ default implementation.
  Readers may override this method to achieve better batched
  I/O performance.
Change: 124012016

1 files changed, 12 insertions, 0 deletions

diff --git a/tensorflow/python/kernel_tests/reader_ops_test.py b/tensorflow/python/kernel_tests/reader_ops_test.py
index 1b55f02285..c49ff3fb49 100644
--- a/tensorflow/python/kernel_tests/reader_ops_test.py
+++ b/tensorflow/python/kernel_tests/reader_ops_test.py
@@ -32,6 +32,18 @@ class IdentityReaderTest(tf.test.TestCase):
     self.assertAllEqual(expected, k)
     self.assertAllEqual(expected, v)
 
+  def testReadUpTo(self):
+    # Note that this just tests the default ReaderReadUpTo
+    # since it is not overriden for IdentityReader.
+    with self.test_session() as sess:
+      reader = tf.IdentityReader("test_reader")
+      queue = tf.FIFOQueue(99, [tf.string], shapes=())
+      keys, values = reader.read_up_to(queue, 3)
+      queue.enqueue_many([["A", "B", "C"]]).run()
+      k, v = sess.run([keys, values])
+      self.assertAllEqual([b"A", b"B", b"C"], k)
+      self.assertAllEqual([b"A", b"B", b"C"], v)
+
   def testOneEpoch(self):
     with self.test_session() as sess:
       reader = tf.IdentityReader("test_reader")