Add to_dataset to IOIterableTensor (e.g., Kafka) as well

Signed-off-by: Yong Tang <[email protected]>

Author: yongtang

tensorflow_io/core/kernels/io_interface.h

@@ -160,11 +160,7 @@ class IOIterableNextOp : public OpKernel {
     OP_REQUIRES_OK(context, resource->Next(capacity, tensors, &record_read));
     for (size_t i = 0; i < tensors.size(); i++) {
       if (record_read < capacity) {
-        gtl::InlinedVector<int64, 4> dims = shapes[i].dim_sizes();
-        dims[0] = record_read;
-        Tensor value;
-        value.CopyFrom(tensors[i], TensorShape(dims));
-        context->set_output(i, value);
+        context->set_output(i, tensors[i].Slice(0, record_read));
       } else {
         context->set_output(i, tensors[i]);
       }

tensorflow_io/core/python/ops/io_tensor.py

@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import sys
 import collections
 
 import tensorflow as tf
@@ -325,7 +326,55 @@ def __init__(self,
   # Iterator
   #=============================================================================
   def __iter__(self):
-    return self._function()
+    resource = self._function["init"](self._function["data"])
+    capacity = 1
+    while True:
+      value = self._function["next"](resource, capacity=capacity)
+      if tf.shape(value)[0].numpy() < capacity:
+        return
+      yield value
+
+  #=============================================================================
+  # Dataset Conversions
+  #=============================================================================
+
+  def to_dataset(self):
+    """Converts this `IOIterableTensor` into a `tf.data.Dataset`.
+
+    Example:
+
+    ```python
+    ```
+
+    Args:
+
+    Returns:
+      A `tf.data.Dataset` with value obtained from this `IOIterableTensor`.
+    """
+    class _IOIterableTensorDataset(data_ops.BaseDataset):
+      """_IOIterableTensorDataset"""
+
+      def __init__(self, dtype, shape, function):
+        func_init = function["init"]
+        func_next = function["next"]
+        func_data = function["data"]
+        resource = func_init(func_data)
+        capacity = 4096
+        dataset = data_ops.BaseDataset.range(
+            0, sys.maxsize, capacity).map(
+                lambda i: func_next(resource, capacity)).apply(
+                    tf.data.experimental.take_while(
+                        lambda v: tf.greater(tf.shape(v)[0], 0))).apply(
+                            tf.data.experimental.unbatch())
+        self._dataset = dataset
+        self._resource = resource
+        self._function = function
+        shape = shape[1:]
+        super(_IOIterableTensorDataset, self).__init__(
+            self._dataset._variant_tensor, [dtype], [shape]) # pylint: disable=protected-access
+
+    return _IOIterableTensorDataset(
+        self._dtype, self._shape, self._function)
 
 class AudioIOTensor(IOTensor):
   """AudioIOTensor"""
@@ -369,30 +418,31 @@ def __init__(self,
                subscription, servers, timeout, eof, conf,
                internal=False):
     with tf.name_scope("KafkaIOTensor") as scope:
-      dtype = tf.string
-      shape = tf.TensorShape([None])
 
-      def function():
-        """_iter_func"""
-        metadata = []
-        if servers is not None:
-          metadata.append("bootstrap.servers=%s" % servers)
-        if timeout is not None:
-          metadata.append("conf.timeout=%d" % timeout)
-        if eof is not None:
-          metadata.append("conf.eof=%d" % (1 if eof else 0))
-        if conf is not None:
-          for e in conf:
-            metadata.append(e)
+      metadata = []
+      if servers is not None:
+        metadata.append("bootstrap.servers=%s" % servers)
+      if timeout is not None:
+        metadata.append("conf.timeout=%d" % timeout)
+      if eof is not None:
+        metadata.append("conf.eof=%d" % (1 if eof else 0))
+      if conf is not None:
+        for e in conf:
+          metadata.append(e)
+
+      func_data = {"subscription": subscription, "metadata": metadata}
+      def func_init(data):
+        """func_init"""
         resource, _, _ = kafka_ops.kafka_iterable_init(
-            subscription, metadata=metadata,
+            data["subscription"], metadata=data["metadata"],
             container=scope, shared_name=subscription)
-        capacity = 1
-        while True:
-          value = kafka_ops.kafka_iterable_next(resource, capacity=capacity)
-          if tf.shape(value)[0].numpy() < capacity:
-            return
-          yield value
+        return resource
+      func_next = kafka_ops.kafka_iterable_next
+
+      dtype = tf.string
+      shape = tf.TensorShape([None])
 
       super(KafkaIOTensor, self).__init__(
-          dtype, shape, function, properties=None, internal=internal)
+          dtype, shape,
+          {"init": func_init, "next": func_next, "data": func_data},
+          properties=None, internal=internal)

tensorflow_io/kafka/kernels/kafka_kernels.cc

@@ -179,12 +179,13 @@ class KafkaIterable : public IOIterableInterface {
   }
   Status Next(const int64 capacity, std::vector<Tensor>& tensors, int64* record_read) override {
     *record_read = 0;
-    while ((*record_read) < capacity) {
+    while (consumer_.get() != nullptr && (*record_read) < capacity) {
       if (!kafka_event_cb_.run()) {
         return errors::Internal("failed to consume due to all brokers down");
       }
       if (range_.second >= 0 && (subscription_->offset() >= range_.second || offset_ >= range_.second)) {
         // EOF of topic
+        consumer_.reset(nullptr);
         return Status::OK();
       }
 
@@ -200,6 +201,7 @@ class KafkaIterable : public IOIterableInterface {
       if (message->err() == RdKafka::ERR__PARTITION_EOF) {
         LOG(INFO) << "Partition reach EOF, current offset: " << offset_;
         if (eof_) {
+          consumer_.reset(nullptr);
           return Status::OK();
         }
       }

tests/test_kafka_eager.py

@@ -34,6 +34,11 @@ def test_kafka_io_tensor():
       e.numpy() for e in kafka] == [
           ("D" + str(i)).encode() for i in range(10)])
 
+  dataset = kafka.to_dataset().batch(2)
+  assert np.all([
+      e.numpy().tolist() for e in dataset] == np.asarray([
+          ("D" + str(i)).encode() for i in range(10)]).reshape((5, 2)))
+
 @pytest.mark.skipif(
     not (hasattr(tf, "version") and
          tf.version.VERSION.startswith("2.0.")), reason=None)