[SPARK-49442][SS] Cache Kafka partition metadata to reduce DescribeTopics RPCs per micro-batch

connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala

@@ -57,6 +57,13 @@ private[kafka010] class KafkaOffsetReaderAdmin(
   private[kafka010] val offsetFetchAttemptIntervalMs =
     readerOptions.getOrElse(KafkaSourceProvider.FETCH_OFFSET_RETRY_INTERVAL_MS, "1000").toLong
 
+  private val partitionMetadataCacheTtlMs: Long =
+    readerOptions.getOrElse(KafkaSourceProvider.PARTITION_METADATA_CACHE_TTL_MS, "-1").toLong
+
+  // Protected by this.synchronized (always accessed inside withRetries, which holds the lock).
+  private var cachedPartitions: Set[TopicPartition] = Set.empty
+  private var cacheTimestampMs: Long = 0L
+
   /**
    * An AdminClient used in the driver to query the latest Kafka offsets.
    * This only queries the offsets because AdminClient has no functionality to commit offsets like
@@ -127,7 +134,7 @@ private[kafka010] class KafkaOffsetReaderAdmin(
       logDebug(s"Assigned partitions: $partitions. Seeking to $partitionOffsets")
       partitionOffsets
     }
-    val partitions = withRetries { consumerStrategy.assignedTopicPartitions(admin) }
+    val partitions = withRetries { resolvePartitions() }
     // Obtain TopicPartition offsets with late binding support
     offsetRangeLimit match {
       case EarliestOffsetRangeLimit => partitions.map {
@@ -439,12 +446,30 @@ private[kafka010] class KafkaOffsetReaderAdmin(
     }
   }
 
+  // Must be called inside withRetries (which holds this.synchronized).
+  private def resolvePartitions(): Set[TopicPartition] = {
+    if (partitionMetadataCacheTtlMs <= 0) {
+      return consumerStrategy.assignedTopicPartitions(admin)
+    }
+    val now = System.currentTimeMillis()
+    if (cacheTimestampMs > 0 && (now - cacheTimestampMs) < partitionMetadataCacheTtlMs) {
+      logDebug(s"Reusing cached partitions (age ${now - cacheTimestampMs}ms < " +
+        s"${partitionMetadataCacheTtlMs}ms TTL): $cachedPartitions")
+      cachedPartitions
+    } else {
+      val fresh = consumerStrategy.assignedTopicPartitions(admin)
+      cachedPartitions = fresh
+      cacheTimestampMs = now
+      fresh
+    }
+  }
+
   private def partitionsAssignedToAdmin(
       body: ju.Set[TopicPartition] => Map[TopicPartition, Long])
     : Map[TopicPartition, Long] = {
 
     withRetries {
-      val partitions = consumerStrategy.assignedTopicPartitions(admin).asJava
+      val partitions = resolvePartitions().asJava
       logDebug(s"Partitions assigned: $partitions.")
       body(partitions)
     }
@@ -493,5 +518,7 @@ private[kafka010] class KafkaOffsetReaderAdmin(
   private def resetAdmin(): Unit = synchronized {
     stopAdmin()
     _admin = null  // will automatically get reinitialized again
+    cachedPartitions = Set.empty
+    cacheTimestampMs = 0L
   }
 }

connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala

@@ -274,6 +274,15 @@ private[kafka010] class KafkaSourceProvider extends DataSourceRegister
       if (p <= 0) throw new IllegalArgumentException("minPartitions must be positive")
     }
 
+    if (params.contains(PARTITION_METADATA_CACHE_TTL_MS)) {
+      val v = params(PARTITION_METADATA_CACHE_TTL_MS).toLong
+      if (v != -1 && v <= 0) {
+        throw new IllegalArgumentException(
+          "Option 'partition.metadata.cache.ttl.ms' must be a positive number of milliseconds, " +
+          "or -1 to disable caching")
+      }
+    }
+
     if (params.contains(MAX_RECORDS_PER_PARTITION_OPTION_KEY)) {
       val p = params(MAX_RECORDS_PER_PARTITION_OPTION_KEY).toLong
       if (p <= 0) {
@@ -577,6 +586,7 @@ private[kafka010] object KafkaSourceProvider extends Logging {
   private[kafka010] val DEFAULT_MAX_TRIGGER_DELAY = "15m"
   private[kafka010] val FETCH_OFFSET_NUM_RETRY = "fetchoffset.numretries"
   private[kafka010] val FETCH_OFFSET_RETRY_INTERVAL_MS = "fetchoffset.retryintervalms"
+  private[kafka010] val PARTITION_METADATA_CACHE_TTL_MS = "partition.metadata.cache.ttl.ms"
   private[kafka010] val CONSUMER_POLL_TIMEOUT = "kafkaconsumer.polltimeoutms"
   private[kafka010] val STARTING_OFFSETS_BY_TIMESTAMP_STRATEGY_KEY =
     "startingoffsetsbytimestampstrategy"

connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderSuite.scala

@@ -323,4 +323,90 @@ class KafkaOffsetReaderSuite extends SharedSparkSession with KafkaTest {
       reader.close()
     }
   }
+
+  // SPARK-49442: When partition.metadata.cache.ttl.ms is set, repeated fetchLatestOffsets
+  // calls within the TTL window must reuse the cached partition set and issue only one
+  // DescribeTopics RPC to the broker, regardless of how many micro-batches run.
+  test("SPARK-49442: partition.metadata.cache.ttl.ms suppresses redundant " +
+      "DescribeTopics RPCs within the TTL window") {
+    val topic = newTopic()
+    testUtils.createTopic(topic, partitions = 3)
+
+    val ttlMs = 300000 // 5 minutes
+    val describeCount = new AtomicInteger(0)
+    val countingStrategy = new SubscribeStrategy(Seq(topic)) {
+      override def assignedTopicPartitions(admin: Admin): Set[TopicPartition] = {
+        describeCount.incrementAndGet()
+        super.assignedTopicPartitions(admin)
+      }
+    }
+
+    val reader = new KafkaOffsetReaderAdmin(
+      countingStrategy,
+      KafkaSourceProvider.kafkaParamsForDriver(Map(
+        "bootstrap.servers" -> testUtils.brokerAddress
+      )),
+      CaseInsensitiveMap(Map(
+        KafkaSourceProvider.PARTITION_METADATA_CACHE_TTL_MS -> ttlMs.toString
+      )),
+      ""
+    )
+
+    try {
+      val numBatches = 5
+      for (_ <- 1 to numBatches) {
+        reader.fetchLatestOffsets(None)
+      }
+      // All 5 fetches fall within the TTL window, so the partition set is resolved only once.
+      assert(describeCount.get() === 1,
+        s"Expected 1 DescribeTopics call but got ${describeCount.get()}")
+    } finally {
+      reader.close()
+    }
+  }
+
+  test("SPARK-49442: partition.metadata.cache.ttl.ms refreshes after TTL expires") {
+    val topic = newTopic()
+    testUtils.createTopic(topic, partitions = 3)
+
+    val ttlMs = 2000 // large enough to not expire during two sequential fetches
+    val describeCount = new AtomicInteger(0)
+    val countingStrategy = new SubscribeStrategy(Seq(topic)) {
+      override def assignedTopicPartitions(admin: Admin): Set[TopicPartition] = {
+        describeCount.incrementAndGet()
+        super.assignedTopicPartitions(admin)
+      }
+    }
+
+    val reader = new KafkaOffsetReaderAdmin(
+      countingStrategy,
+      KafkaSourceProvider.kafkaParamsForDriver(Map(
+        "bootstrap.servers" -> testUtils.brokerAddress
+      )),
+      CaseInsensitiveMap(Map(
+        KafkaSourceProvider.PARTITION_METADATA_CACHE_TTL_MS -> ttlMs.toString
+      )),
+      ""
+    )
+
+    try {
+      // First fetch populates the cache (1 RPC).
+      reader.fetchLatestOffsets(None)
+      assert(describeCount.get() === 1)
+
+      // Fetches within the TTL reuse the cache (still 1 RPC).
+      reader.fetchLatestOffsets(None)
+      assert(describeCount.get() === 1)
+
+      // Sleep past the TTL so the next fetch must refresh.
+      Thread.sleep(ttlMs + 50)
+
+      // Cache has expired: one more RPC expected.
+      reader.fetchLatestOffsets(None)
+      assert(describeCount.get() === 2,
+        s"Expected 2 DescribeTopics calls after TTL expiry but got ${describeCount.get()}")
+    } finally {
+      reader.close()
+    }
+  }
 }

connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaSourceProviderSuite.scala

@@ -133,4 +133,25 @@ class KafkaSourceProviderSuite extends SparkFunSuite {
     val provider = new KafkaSourceProvider()
     provider.getTable(options).newScanBuilder(options).build()
   }
+
+  test("SPARK-49442: partition.metadata.cache.ttl.ms validation") {
+    val sparkEnv = mock(classOf[SparkEnv])
+    when(sparkEnv.conf).thenReturn(new SparkConf())
+    SparkEnv.set(sparkEnv)
+
+    // -1 (disabled) and positive values are valid; validation fires inside toMicroBatchStream
+    Seq("-1", "1", "30000").foreach { v =>
+      val options = buildKafkaSourceCaseInsensitiveStringMap(
+        KafkaSourceProvider.PARTITION_METADATA_CACHE_TTL_MS -> v)
+      getKafkaDataSourceScan(options).toMicroBatchStream("dummy")
+    }
+    // 0 and other non-positive values (except -1) are invalid
+    Seq("0", "-2", "-100").foreach { v =>
+      val options = buildKafkaSourceCaseInsensitiveStringMap(
+        KafkaSourceProvider.PARTITION_METADATA_CACHE_TTL_MS -> v)
+      intercept[IllegalArgumentException] {
+        getKafkaDataSourceScan(options).toMicroBatchStream("dummy")
+      }
+    }
+  }
 }

docs/streaming/structured-streaming-kafka-integration.md

@@ -481,6 +481,19 @@ The following configurations are optional:
   <td>streaming and batch</td>
   <td>milliseconds to wait before retrying to fetch Kafka offsets</td>
 </tr>
+<tr>
+  <td>partition.metadata.cache.ttl.ms</td>
+  <td>long</td>
+  <td>-1</td>
+  <td>streaming and batch</td>
+  <td>How long (in milliseconds) to cache the set of assigned Kafka topic-partitions in the driver
+  before issuing a fresh <code>DescribeTopics</code> RPC to the broker. When set to a positive value,
+  repeated offset fetches within the TTL window reuse the cached partition set instead of querying
+  the broker each time, reducing metadata load on the broker and on the driver. This mirrors
+  the semantics of the Kafka client's own <code>metadata.max.age.ms</code> setting. Set to -1
+  (the default) to disable caching (existing behavior). Note: newly added partitions will not be
+  discovered until the cache expires.</td>
+</tr>
 <tr>
   <td>maxOffsetsPerTrigger</td>
   <td>long</td>