Removed tfio.Dataset exposure, instead added tfio.IOTensor.to_dataset()

Instead of exposing a tfio.Dataset, I have added a `tfio.IOTensor.to_dataset()`
which I think makes much more sense to avoid the confusion.

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

Author: yongtang

tensorflow_io/__init__.py

@@ -18,4 +18,3 @@
 from __future__ import print_function
 
 from tensorflow_io.core.python.ops.io_tensor import IOTensor
-from tensorflow_io.core.python.ops.dataset_ops import Dataset

tensorflow_io/core/python/ops/dataset_ops.py

@@ -21,6 +21,7 @@
 
 import tensorflow as tf
 from tensorflow_io.core.python.ops import core_ops
+from tensorflow_io.core.python.ops import data_ops
 from tensorflow_io.kafka.python.ops.kafka_ops import kafka_ops
 
 class _IOBaseTensor(object):
@@ -261,6 +262,48 @@ def to_tensor(self, **kwargs):
     with tf.name_scope(kwargs.get("name", "IOToTensor")):
       return self.__getitem__(slice(None, None))
 
+  #=============================================================================
+  # Dataset Conversions
+  #=============================================================================
+
+  def to_dataset(self):
+    """Converts this `IOTensor` into a `tf.data.Dataset`.
+
+    Example:
+
+    ```python
+    ```
+
+    Args:
+
+    Returns:
+      A `tf.data.Dataset` with value obtained from this `IOTensor`.
+    """
+    class _IOTensorDataset(data_ops.BaseDataset):
+      """_IOTensorDataset"""
+
+      def __init__(self, dtype, shape, resource, function):
+        start = 0
+        stop = shape[0]
+        capacity = 4096
+        entry_start = list(range(start, stop, capacity))
+        entry_stop = entry_start[1:] + [stop]
+        dataset = data_ops.BaseDataset.from_tensor_slices((
+            tf.constant(entry_start, tf.int64),
+            tf.constant(entry_stop, tf.int64))).map(
+                lambda start, stop: function(
+                    resource, start, stop, 1, dtype=dtype)).apply(
+                        tf.data.experimental.unbatch())
+        self._dataset = dataset
+        self._resource = resource
+        self._function = function
+        shape = shape[1:]
+        super(_IOTensorDataset, self).__init__(
+            self._dataset._variant_tensor, [dtype], [shape]) # pylint: disable=protected-access
+
+    return _IOTensorDataset(
+        self._dtype, self._shape, self._resource, self._function)
+
 class IOIterableTensor(_IOBaseTensor):
   """IOIterableTensor"""
 

tensorflow_io/core/python/ops/io_tensor.py

@@ -24,6 +24,7 @@
 if not (hasattr(tf, "version") and tf.version.VERSION.startswith("2.")):
   tf.compat.v1.enable_eager_execution()
 import tensorflow_io as tfio # pylint: disable=wrong-import-position
+import tensorflow_io.audio as audio_io # pylint: disable=wrong-import-position
 
 audio_path = os.path.join(
     os.path.dirname(os.path.abspath(__file__)),
@@ -37,14 +38,14 @@ def test_audio_dataset():
 
   f = lambda x: float(x) / (1 << 15)
 
-  audio_dataset = tfio.Dataset.from_audio(audio_path)
+  audio_dataset = audio_io.WAVDataset(audio_path)
   i = 0
   for v in audio_dataset:
     assert audio_v.audio[i].numpy() == f(v.numpy())
     i += 1
   assert i == 5760
 
-  audio_dataset = tfio.Dataset.from_audio(audio_path).batch(2)
+  audio_dataset = tfio.IOTensor.from_audio(audio_path).to_dataset().batch(2)
   i = 0
   for v in audio_dataset:
     assert audio_v.audio[i].numpy() == f(v[0].numpy())