diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala index 69f346af19a88..f182e156c76e9 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala +++ b/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 } } diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala index ff2f16d26b932..b5a7bcb98fba2 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala +++ b/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" diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderSuite.scala index 341ad14a4fbb9..bed76cb8990a6 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderSuite.scala +++ b/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() + } + } } diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaSourceProviderSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaSourceProviderSuite.scala index 303394b445792..2601b118132fa 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaSourceProviderSuite.scala +++ b/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") + } + } + } } diff --git a/docs/streaming/structured-streaming-kafka-integration.md b/docs/streaming/structured-streaming-kafka-integration.md index 6cefcb561f224..364fc77f1ef1b 100644 --- a/docs/streaming/structured-streaming-kafka-integration.md +++ b/docs/streaming/structured-streaming-kafka-integration.md @@ -481,6 +481,19 @@ The following configurations are optional: streaming and batch milliseconds to wait before retrying to fetch Kafka offsets + + partition.metadata.cache.ttl.ms + long + -1 + streaming and batch + How long (in milliseconds) to cache the set of assigned Kafka topic-partitions in the driver + before issuing a fresh DescribeTopics 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 metadata.max.age.ms setting. Set to -1 + (the default) to disable caching (existing behavior). Note: newly added partitions will not be + discovered until the cache expires. + maxOffsetsPerTrigger long