This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c369cf28e26d [SPARK-49442][SS] Cache Kafka partition metadata to
reduce DescribeTopics RPCs per micro-batch
c369cf28e26d is described below
commit c369cf28e26d3c941e177a0e894308dbef58bbeb
Author: Vinod KC <[email protected]>
AuthorDate: Tue Jul 21 07:55:08 2026 +0900
[SPARK-49442][SS] Cache Kafka partition metadata to reduce DescribeTopics
RPCs per micro-batch
### What changes were proposed in this pull request?
Adds a new source option `partition.metadata.cache.ttl.ms` to
`KafkaOffsetReaderAdmin`. When set to a positive value, the driver caches the
set of assigned `TopicPartition`s and reuses it for all offset fetch calls
within the TTL window, skipping the `DescribeTopics` RPC to the broker. The
cache is invalidated on admin reset (i.e., on retry). Default is `-1` (caching
disabled, preserving existing behavior).
### Why are the changes needed?
Every micro-batch unconditionally called
`consumerStrategy.assignedTopicPartitions(admin)`, which issues a
`DescribeTopics` RPC to the Kafka broker even though partition assignments are
stable in most production deployments. For deployments with large partition
counts and short trigger intervals, this generates unnecessary metadata traffic.
### Does this PR introduce _any_ user-facing change?
Yes. Adds a new source option `partition.metadata.cache.ttl.ms` (long,
default `-1`). Valid values are `-1` (caching disabled) or any positive integer
representing milliseconds. Documentation is updated accordingly. Existing
behavior is unchanged when the option is not set.
> **Note:** when caching is enabled, newly added partitions will not be
discovered until the cache expires.
### How was this patch tested?
Added tests in `KafkaOffsetReaderSuite` `KafkaSourceProviderSuite`,
`KafkaOffsetReaderSuite` and `KafkaSourceProviderSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Yes, 'Generated-by: Claude (Sonnet 4.6)
Closes #57351 from vinodkc/spark-49442.
Authored-by: Vinod KC <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../sql/kafka010/KafkaOffsetReaderAdmin.scala | 31 +++++++-
.../spark/sql/kafka010/KafkaSourceProvider.scala | 10 +++
.../sql/kafka010/KafkaOffsetReaderSuite.scala | 86 ++++++++++++++++++++++
.../sql/kafka010/KafkaSourceProviderSuite.scala | 21 ++++++
.../structured-streaming-kafka-integration.md | 13 ++++
5 files changed, 159 insertions(+), 2 deletions(-)
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 69f346af19a8..f182e156c76e 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 ff2f16d26b93..b5a7bcb98fba 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 341ad14a4fbb..bed76cb8990a 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 303394b44579..2601b118132f 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 6cefcb561f22..364fc77f1ef1 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:
<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>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
