message: Rework on tensorflow_io/kafka/python/ops/kafka_dataset_ops.py
diff: @@ -17,13 +17,16 @@ from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  -from tensorflow.contrib.kafka.python.ops import gen_dataset_ops
  -from tensorflow.contrib.kafka.python.ops import kafka_op_loader # pylint: disable=unused-import
  +import os
  +
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  +from tensorflow.python.framework import load_library
  +kafka_ops = load_library.load_op_library(
  + os.path.join("tensorflow_io", "kafka", "python", "ops", "_kafka_ops.so"))
  class KafkaDataset(dataset_ops.DatasetSource):
  """A Kafka Dataset that consumes the message.
  @@ -59,7 +62,7 @@ class KafkaDataset(dataset_ops.DatasetSource):
  timeout, dtype=dtypes.int64, name="timeout")
  def _as_variant_tensor(self):
  - return gen_dataset_ops.kafka_dataset(self._topics, self._servers,
  + return kafka_ops.kafka_dataset(self._topics, self._servers,
  self._group, self._eof, self._timeout)
  @property
  