[KafkaDataset] Support KafkaDataset. (DeepRec-AI#118)

WORKSPACE

@@ -252,6 +252,22 @@ http_archive(
     urls = ["https://github.com/cameron314/readerwriterqueue/archive/v1.0.1.tar.gz"],
 )
 
+http_archive(
+    name = "kafka150",
+    build_file = "//third_party:kafka.BUILD",
+    patch_cmds = [
+        "rm -f src/win32_config.h",
+        # TODO: Remove the fowllowing once librdkafka issue is resolved.
+        """sed -i.bak '\|rd_kafka_log(rk,|,/ exceeded);/ s/^/\/\//' src/rdkafka_cgrp.c""",
+    ],
+    sha256 = "f7fee59fdbf1286ec23ef0b35b2dfb41031c8727c90ced6435b8cf576f23a656",
+    strip_prefix = "librdkafka-1.5.0",
+    urls = [
+        "https://mirror.tensorflow.org/github.com/edenhill/librdkafka/archive/v1.5.0.tar.gz",
+        "https://github.com/edenhill/librdkafka/archive/v1.5.0.tar.gz",
+    ],
+)
+
 load("//third_party:repo.bzl", "tf_http_archive")
 tf_http_archive(
     name = "seastar_repo",

tensorflow/core/kernels/data/BUILD

@@ -1075,6 +1075,20 @@ tf_cc_test(
     ],
 )
 
+tf_kernel_library(
+    name = "kafka_dataset_op",
+    srcs = ["kafka_dataset_op.cc"],
+    deps = [
+        "@kafka150",
+        ":name_utils",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:dataset_ops_op_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+    ],
+)
+
 tf_kernel_library(
     name = "iterator_ops",
     srcs = ["iterator_ops.cc"],
@@ -1280,6 +1294,7 @@ tf_kernel_library(
         ":tensor_slice_dataset_op",
         ":text_line_dataset_op",
         ":tf_record_dataset_op",
+        ":kafka_dataset_op",
         ":window_dataset_op",
         ":zip_dataset_op",
         "//tensorflow/core:array_ops_op_lib",