fixed a bug related to retrieving offsets (#1063)

Author: kvignesh1420

tensorflow_io/core/kernels/kafka_kernels.cc

@@ -772,8 +772,15 @@ class KafkaRebalanceCb : public RdKafka::RebalanceCb {
   void rebalance_cb(RdKafka::KafkaConsumer* consumer, RdKafka::ErrorCode err,
                     std::vector<RdKafka::TopicPartition*>& partitions) {
     LOG(ERROR) << "REBALANCE: " << RdKafka::err2str(err);
+    LOG(ERROR) << "Retrieved committed offsets with status code: "
+               << consumer->committed(partitions, 5000);
+
     for (int partition = 0; partition < partitions.size(); partition++) {
-      partitions[partition]->set_offset(RdKafka::Topic::OFFSET_STORED);
+      if (partitions[partition]->offset() == -1001) {
+        LOG(INFO)
+            << "The consumer group was newly created, reading from beginning";
+        partitions[partition]->set_offset(RdKafka::Topic::OFFSET_BEGINNING);
+      }
       LOG(INFO) << "REBALANCE: " << partitions[partition]->topic() << "["
                 << partitions[partition]->partition() << "], "
                 << partitions[partition]->offset() << " "

tests/test_kafka_eager.py

@@ -184,3 +184,32 @@ def test_avro_encode_decode():
     entries = tfio.experimental.serialization.decode_avro(message, schema=schema)
     assert np.all(entries["f1"].numpy() == f1.numpy())
     assert np.all(entries["f2"].numpy() == f2.numpy())
+
+
+def test_kafka_group_io_dataset_new_cg():
+    """Test the functionality of the KafkaGroupIODataset when the consumer group
+    is being newly created.
+    """
+    dataset = tfio.experimental.streaming.KafkaGroupIODataset(
+        topics=["key-partition-test"],
+        group_id="cgtest",
+        servers="localhost:9092",
+        configuration=["session.timeout.ms=7000", "max.poll.interval.ms=8000"],
+    )
+    assert np.all(
+        sorted([k.numpy() for (k, _) in dataset])
+        == sorted([("D" + str(i % 10)).encode() for i in range(10)])
+    )
+
+
+def test_kafka_group_io_dataset_no_lag():
+    """Test the functionality of the KafkaGroupIODataset when the
+    consumer group has read all the messages and committed the offsets.
+    """
+    dataset = tfio.experimental.streaming.KafkaGroupIODataset(
+        topics=["key-partition-test"],
+        group_id="cgtest",
+        servers="localhost:9092",
+        configuration=["session.timeout.ms=7000", "max.poll.interval.ms=8000"],
+    )
+    assert np.all(sorted([k.numpy() for (k, _) in dataset]) == [])