Bug: KafkaBatchPartitionReader blocks until timeout when transactional messages are enabled
Status: Open | Priority: Major | Resolution: 
Created: 27/May/22 05:14
Affects Version/s: 3.2.1

Steps to reproduce:

 

1. Consumer structured stream:

 
{code:java}
val df = spark
  .readStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "host1:port1,host2:port2")
  .option("kafka.isolation.level", "read_committed")
  .option("subscribe", "topic1")
  .load()
df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)")
  .as[(String, String)]
  .writeStream
  .foreachBatch { (batchDF: DataFrame, batchId: Long) =>     
     batchDF.cache()
     batchDF.write.format("console").mode(SaveMode.Append).save()
     batchDF.unpersist()
   }.start(){code}
2. Producer
{code:java}
val producer = new KafkaProducer[String, String](...) 
producer.initTransactions()
try {
  producer.beginTransaction()
  producer.send(new ProducerRecord("topic1", "key", "value")).get()
  producer.commitTransaction()
} finally {
  producer.close()
}{code}
Expected: spark successfully processes the record and writes it to the output sink as soon as it's read from kafka

Actual: the stream gets stuck for 2 minutes (default kafka poll timeout) 

 

This is a really tricky one. It boils down to the fact that org.apache.spark.sql.kafka010.KafkaBatchPartitionReader when called from org.apache.spark.sql.execution.datasources.v2.PartitionIterator does not follow iterator semantics correctly. It calls org.apache.spark.sql.kafka010.consumer.KafkaDataConsumer#get the first time it encounters an offset with an 'invisible' kafka message and returns null, however if PartitionIterator.hasNext() is called twice by something up the stack (in this particular case, by a DefaultCachedBatchSerializer), the KafkaBatchPartitionReader chokes and does not handle the 'invisible' message correctly the second time around, instead trying to call poll() again and blocking until timeout.

 

As far as i can see, there is not workaround if you need to use foreachBatch

 

Example log output:
{noformat}
groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Subscribed to partition(s): topic1-43
2022-05-26 16:50:59,189 DEBUG consumer.KafkaDataConsumer: Get spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor topic1-43 nextOffset -2 requested 1470527
2022-05-26 16:50:59,189 DEBUG consumer.InternalKafkaConsumer: Seeking to spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor topic1-43 1470527
2022-05-26 16:50:59,189 INFO consumer.KafkaConsumer: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Seeking to offset 1470527 for partition topic1-43
2022-05-26 16:50:59,235 INFO clients.Metadata: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Cluster ID: N0b937OTSdCNOpAvSrTDVw
2022-05-26 16:50:59,314 DEBUG consumer.InternalKafkaConsumer: Polled spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor [topic1-43]  1
2022-05-26 16:50:59,314 DEBUG consumer.InternalKafkaConsumer: Offset changed from 1470527 to 1470529 after polling
2022-05-26 16:50:59,314 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Seeking to EARLIEST offset of partition topic1-43
2022-05-26 16:50:59,854 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Resetting offset for partition topic1-43 to position FetchPosition{offset=1465967, offsetEpoch=Optional.empty, currentLeader=LeaderAndEpoch{leader=Optional[loninkafkuu6.uk.db.com:9093 (id: 1000000006 rack: null)], epoch=1450}}.
2022-05-26 16:50:59,854 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Seeking to LATEST offset of partition topic1-43
2022-05-26 16:50:59,863 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Resetting offset for partition topic1-43 to position FetchPosition{offset=1470529, offsetEpoch=Optional.empty, currentLeader=LeaderAndEpoch{leader=Optional[loninkafkuu6.uk.db.com:9093 (id: 1000000006 rack: null)], epoch=1450}}.
2022-05-26 16:50:59,865 DEBUG consumer.KafkaDataConsumer: Get spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor topic1-43 nextOffset 1470528 requested 1470528
2022-05-26 16:50:59,866 DEBUG consumer.KafkaDataConsumer: Get spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor topic1-43 nextOffset -2 requested 1470528
2022-05-26 16:50:59,866 DEBUG consumer.InternalKafkaConsumer: Seeking to spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor topic1-43 1470528
2022-05-26 16:50:59,866 INFO consumer.KafkaConsumer: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Seeking to offset 1470528 for partition topic1-43
2022-05-26 16:52:59,866 DEBUG consumer.InternalKafkaConsumer: Polled spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor []  0
2022-05-26 16:52:59,867 DEBUG consumer.InternalKafkaConsumer: Offset changed from 1470528 to 1470529 after polling
2022-05-26 16:52:59,867 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Seeking to EARLIEST offset of partition topic1-43
2022-05-26 16:53:00,046 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Resetting offset for partition topic1-43 to position FetchPosition{offset=1465967, offsetEpoch=Optional.empty, currentLeader=LeaderAndEpoch{leader=Optional[loninkafkuu6.uk.db.com:9093 (id: 1000000006 rack: null)], epoch=1450}}.
2022-05-26 16:53:00,046 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Seeking to LATEST offset of partition topic1-43
2022-05-26 16:53:00,053 INFO internals.SubscriptionState: [Consumer clientId=consumer-spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor-452, groupId=spark-kafka-source-30039f5f-9d67-4991-a01f-43ec1f2ad213--1217354974-executor] Resetting offset for partition topic1-43 to position FetchPosition{offset=1470529, offsetEpoch=Optional.empty, currentLeader=LeaderAndEpoch{leader=Optional[loninkafkuu6.uk.db.com:9093 (id: 1000000006 rack: null)], epoch=1450}}.
2022-05-26 16:53:00,055 INFO memory.MemoryStore: Block rdd_6514_0 stored as values in memory (estimated size 3.9 KiB, free 366.2 MiB){noformat}