This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 525f04570b0f [SPARK-58774][SS] Support topic-level earliest/latest
values in the Kafka source startingOffsets/endingOffsets JSON
525f04570b0f is described below
commit 525f04570b0f836719c4a0dc6c00aca0a1399e6c
Author: Joaquín Antonio De Vicente López <[email protected]>
AuthorDate: Tue Aug 25 13:10:03 2026 +0900
[SPARK-58774][SS] Support topic-level earliest/latest values in the Kafka
source startingOffsets/endingOffsets JSON
### What changes were proposed in this pull request?
This PR lets a topic key in the `startingOffsets` / `endingOffsets` JSON
map to the string `"earliest"` or `"latest"`, in addition to the existing
partition-to-offset object. Both forms can be mixed in the same document:
```
startingOffsets={"topicA":"earliest","topicB":{"0":23,"1":-1},"topicC":"latest"}
```
A topic-level `"earliest"` is equivalent to writing `-2` for every
partition of that topic assigned to the query, and `"latest"` to `-1`. Because
the expansion happens against the discovered partition set rather than a
hard-coded list, no partition count is baked into the option.
Implementation:
- `JsonUtils.specificOffsets` (new) walks the JSON AST and splits
topic-level strings from partition maps. `JsonUtils.partitionOffsets` is left
untouched and stays strict, so checkpointed offsets cannot accidentally accept
the new syntax.
- `SpecificOffsetRangeLimit` gains `topicOffsets: Map[String, Long]` plus a
`resolve(assignedPartitions)` method. The parameter is by-name, so a fully
enumerated limit never pays for discovering the partitions.
- `KafkaOffsetReader.fetchSpecificOffsets` now takes the limit instead of a
bare map; both the admin- and consumer-based readers expand against their
assigned partitions. `fetchPartitionOffsets` reuses the partition set it has
already fetched.
- Batch queries reject a topic bound to `latest` (start) or `earliest`
(end), mirroring the existing per-partition checks.
- A topic-level entry naming a topic with no assigned partition fails with
a new error condition, `KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED`, rather
than silently expanding to nothing, so a typo in a topic name cannot cause the
query to read less than intended.
- `checkOffsetLimitValidity` no longer requires matching topic-partition
key sets between start and end when topic-level entries are present, since the
two sides are legitimately not comparable before resolution. Matching them up
is left to the offset readers, which see the assigned partitions.
JIRA: https://issues.apache.org/jira/browse/SPARK-58774
### Why are the changes needed?
Today the source accepts `"earliest"`, `"latest"`, or a JSON object
enumerating every topic-partition. There is no way to express "earliest for
topic A, latest for topic B" without the fully enumerated form.
In steady state this rarely matters, since a single `"earliest"` or
`"latest"` covers the whole subscription. It becomes acute during incident
recovery or a targeted replay, when one topic has to be rewound while the
others stay where they are. With 24 partitions across 5 subscribed topics, that
is a single-line string of several thousand characters containing 120 entries
whose only content is the sentinel `-2`:
- **Verbosity.** In `.properties` or YAML configuration such a value is a
frequent source of accidental line breaks, which truncate it and surface as a
confusing `IllegalArgumentException`.
- **Opacity.** `-1` and `-2` are magic numbers; a single mistyped sentinel
silently changes the semantics for one partition.
- **Brittleness under repartitioning.** The JSON hard-codes the partition
count. Expanding a topic from 24 to 32 partitions makes the enumerated set stop
matching the assigned one, and the query fails to start with
`KAFKA_START_OFFSET_DOES_NOT_MATCH_ASSIGNED`. Failing loudly is the right
behaviour, but the only remedy is to hand-edit the multi-thousand-character
string after every repartitioning.
- **No composability.** Wanting a different position for a single topic
forces explicit enumeration of all the others.
### Does this PR introduce _any_ user-facing change?
Yes, and it is purely additive.
The topic-level value is disambiguated by JSON type: an object continues to
be parsed as a partition-to-offset map, while a string is the new shorthand.
Every currently valid configuration keeps its meaning, and no previously
accepted document changes behaviour.
Previously:
```
startingOffsets={"topicA":"earliest"}
-> IllegalArgumentException: Expected e.g.
{"topicA":{"0":23,"1":-1},"topicB":{"0":-2}}
```
Now the same option is accepted and applies `earliest` to every partition
of `topicA`.
Two new failure modes, both for input that was previously rejected outright:
- An unrecognised string value (for example `{"topicA":"first"}`) is
rejected at parse time with a message naming the accepted forms.
- A topic-level entry for a topic with no assigned partition raises
`KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED`.
As with the enumerated form, the JSON has to account for every topic
subscribed at the moment the offsets are resolved, so with `subscribePattern`
the matched topics must be known then. Topics matched later are discovered as
new partitions and start at earliest, exactly as today. The documentation for
both options has been updated accordingly.
### How was this patch tested?
New tests:
- `KafkaOffsetRangeLimitSuite` (new file) covers `resolve`: expansion
against the assigned partitions, the short-circuit that avoids fetching
partitions when nothing needs expanding, and the unmatched-topic error.
- `JsonUtilsSuite` covers parsing of the topic-level form, the mixed form,
case insensitivity, and rejection of malformed input.
- `KafkaRelationSuite` covers batch end to end: topic-level offsets, the
mixed form, `subscribePattern`, repartitioning between runs, and the two ways
the offsets can fail to line up with the topics a pattern matches.
- `KafkaMicroBatchSourceSuite` covers streaming end to end: a query started
from topic-level offsets with per-topic strategies, and a topic created after
the initial offsets were resolved being picked up as a new partition rather
than tripping the strict topic check.
- `KafkaRelationSuite` "bad batch query options" gains the topic-level
`latest`/`earliest` rejections.
The batch and streaming tests were verified on both V1 and V2, and the
streaming ones on both the admin- and consumer-based offset readers.
The full `sql-kafka-0-10` suite passes: 37 suites, 602 tests, 0 failures.
`scalastyle` is clean for both `Compile` and `Test`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5)
Closes #58043 from Kino1994/feature/SPARK-58774.
Authored-by: Joaquín Antonio De Vicente López <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
(cherry picked from commit 64b4fd2a72140e7ac43e3d974b908016f8f68148)
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../resources/error/kafka-error-conditions.json | 7 ++
.../org/apache/spark/sql/kafka010/JsonUtils.scala | 49 +++++++++-
.../spark/sql/kafka010/KafkaContinuousStream.scala | 2 +-
.../spark/sql/kafka010/KafkaExceptions.scala | 10 ++
.../spark/sql/kafka010/KafkaMicroBatchStream.scala | 4 +-
.../spark/sql/kafka010/KafkaOffsetRangeLimit.scala | 30 +++++-
.../spark/sql/kafka010/KafkaOffsetReader.scala | 12 ++-
.../sql/kafka010/KafkaOffsetReaderAdmin.scala | 24 +++--
.../sql/kafka010/KafkaOffsetReaderConsumer.scala | 28 +++---
.../apache/spark/sql/kafka010/KafkaSource.scala | 2 +-
.../spark/sql/kafka010/KafkaSourceProvider.scala | 28 +++---
.../apache/spark/sql/kafka010/JsonUtilsSuite.scala | 55 +++++++++++
.../sql/kafka010/KafkaMicroBatchSourceSuite.scala | 90 ++++++++++++++++-
.../sql/kafka010/KafkaOffsetRangeLimitSuite.scala | 58 +++++++++++
.../spark/sql/kafka010/KafkaRelationSuite.scala | 108 ++++++++++++++++++++-
.../structured-streaming-kafka-integration.md | 21 +++-
16 files changed, 483 insertions(+), 45 deletions(-)
diff --git
a/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json
b/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json
index c256b8cbebb2..6aaf067cd3f4 100644
---
a/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json
+++
b/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json
@@ -36,6 +36,13 @@
"Specified: <specifiedPartitions> Assigned: <assignedPartitions>"
]
},
+ "KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED" : {
+ "message" : [
+ "Topics specified with a topic-level offset for Kafka offsets don't have
any assigned partition. Maybe the topics are misspelled, ",
+ "or they are not part of the topics being subscribed.",
+ "Specified: <specifiedTopics> Assigned: <assignedTopics>"
+ ]
+ },
"KAFKA_TIMESTAMP_OFFSET_DOES_NOT_MATCH_ASSIGNED" : {
"message" : [
"Partitions specified for Kafka timestamp based <position> offsets don't
match what are assigned. Maybe topic partitions are created ",
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala
index 7f3eb0370a07..618064d88314 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala
@@ -17,11 +17,14 @@
package org.apache.spark.sql.kafka010
+import java.util.Locale
+
import scala.collection.mutable.HashMap
import scala.util.control.NonFatal
import org.apache.kafka.common.TopicPartition
-import org.json4s.{Formats, NoTypeHints}
+import org.json4s.{Formats, JObject, JString, NoTypeHints}
+import org.json4s.jackson.JsonMethods.parse
import org.json4s.jackson.Serialization
/**
@@ -76,6 +79,50 @@ private object JsonUtils {
}
}
+ /**
+ * Read the offsets of the `startingOffsets` / `endingOffsets` json string.
On top of the
+ * per-TopicPartition form read by [[partitionOffsets]], a topic may bind to
"earliest" or
+ * "latest" as a whole, e.g. {"topicA":"earliest","topicB":{"0":23,"1":-1}}.
Such topic-level
+ * values are returned unexpanded in
`SpecificOffsetRangeLimit.topicOffsets`, since the
+ * partitions of the topic are only known once the offsets get resolved
against Kafka.
+ */
+ def specificOffsets(str: String): SpecificOffsetRangeLimit = {
+ def fail(): Nothing = throw new IllegalArgumentException(
+ s"""Expected e.g. {"topicA":{"0":23,"1":-1},"topicB":{"0":-2}} or
+ |{"topicA":"earliest","topicB":"latest"}, got $str""".stripMargin)
+
+ val partitionOffsets = new HashMap[TopicPartition, Long]
+ val topicOffsets = new HashMap[String, Long]
+ try {
+ parse(str) match {
+ case JObject(topics) =>
+ topics.foreach {
+ case (topic, JString(value)) =>
+ topicOffsets += topic -> (value.toLowerCase(Locale.ROOT) match {
+ case "earliest" => KafkaOffsetRangeLimit.EARLIEST
+ case "latest" => KafkaOffsetRangeLimit.LATEST
+ case _ => fail()
+ })
+ case (topic, partOffsets) =>
+ partOffsets.extract[Map[Int, Long]].foreach { case (part,
offset) =>
+ partitionOffsets += new TopicPartition(topic, part) -> offset
+ }
+ }
+ case _ => fail()
+ }
+ } catch {
+ case NonFatal(_) => fail()
+ }
+ val bothForms =
topicOffsets.keySet.intersect(partitionOffsets.keySet.map(_.topic))
+ if (bothForms.nonEmpty) {
+ throw new IllegalArgumentException(
+ s"""Topic(s) ${bothForms.toSeq.sorted.mkString(", ")} are given both a
topic-level offset
+ |and per-partition offsets, only one of the two forms is allowed
per topic,
+ |got $str""".stripMargin)
+ }
+ SpecificOffsetRangeLimit(partitionOffsets.toMap, topicOffsets.toMap)
+ }
+
def partitionTimestamps(str: String): Map[TopicPartition, Long] = {
try {
Serialization.read[Map[String, Map[Int, Long]]](str).flatMap { case
(topic, partTimestamps) =>
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala
index 041fe074f7e4..e0be21dadc2c 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala
@@ -70,7 +70,7 @@ class KafkaContinuousStream(
val offsets = initialOffsets match {
case EarliestOffsetRangeLimit =>
KafkaSourceOffset(offsetReader.fetchEarliestOffsets())
case LatestOffsetRangeLimit =>
KafkaSourceOffset(offsetReader.fetchLatestOffsets(None))
- case SpecificOffsetRangeLimit(p) => offsetReader.fetchSpecificOffsets(p,
reportDataLoss)
+ case o: SpecificOffsetRangeLimit => offsetReader.fetchSpecificOffsets(o,
reportDataLoss)
case SpecificTimestampRangeLimit(p, strategy) =>
offsetReader.fetchSpecificTimestampBasedOffsets(p, isStartingOffsets =
true, strategy)
case GlobalTimestampRangeLimit(ts, strategy) =>
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala
index 17472fb920f1..4b7e60205a11 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala
@@ -166,6 +166,16 @@ object KafkaExceptions {
"assignedPartitions" -> assignedPartitions.toString))
}
+ def topicOffsetDoesNotMatchAssigned(
+ specifiedTopics: Set[String],
+ assignedTopics: Set[String]): KafkaIllegalStateException = {
+ new KafkaIllegalStateException(
+ errorClass = "KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED",
+ messageParameters = Map(
+ "specifiedTopics" -> specifiedTopics.toString,
+ "assignedTopics" -> assignedTopics.toString))
+ }
+
def timestampOffsetDoesNotMatchAssigned(
isStartingOffsets: Boolean,
specifiedPartitions: Set[TopicPartition],
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala
index 17323165b451..e94c08cd2b18 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala
@@ -379,8 +379,8 @@ private[kafka010] class KafkaMicroBatchStream(
KafkaSourceOffset(kafkaOffsetReader.fetchEarliestOffsets())
case LatestOffsetRangeLimit =>
KafkaSourceOffset(kafkaOffsetReader.fetchLatestOffsets(None))
- case SpecificOffsetRangeLimit(p) =>
- kafkaOffsetReader.fetchSpecificOffsets(p, reportDataLoss)
+ case o: SpecificOffsetRangeLimit =>
+ kafkaOffsetReader.fetchSpecificOffsets(o, reportDataLoss)
case SpecificTimestampRangeLimit(p, strategy) =>
kafkaOffsetReader.fetchSpecificTimestampBasedOffsets(p,
isStartingOffsets = true, strategy)
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala
index ef2feb679afa..365e8cb434b9 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala
@@ -40,9 +40,37 @@ private[kafka010] case object LatestOffsetRangeLimit extends
KafkaOffsetRangeLim
/**
* Represents the desire to bind to specific offsets. A offset == -1 binds to
the
* latest offset, and offset == -2 binds to the earliest offset.
+ *
+ * `topicOffsets` holds topic-level bindings, written in the JSON as
"earliest" or "latest" in
+ * place of a topic's per-partition object. They are expanded against the
partitions discovered
+ * for the topic when the offsets are resolved, so they keep working across
repartitioning.
*/
private[kafka010] case class SpecificOffsetRangeLimit(
- partitionOffsets: Map[TopicPartition, Long]) extends KafkaOffsetRangeLimit
+ partitionOffsets: Map[TopicPartition, Long],
+ topicOffsets: Map[String, Long] = Map.empty) extends KafkaOffsetRangeLimit
{
+
+ /**
+ * Expands the topic-level bindings against `assignedPartitions` and returns
the resulting
+ * offset per topic-partition. `assignedPartitions` is by-name so that a
fully enumerated
+ * limit never pays for discovering the partitions.
+ */
+ def resolve(assignedPartitions: => Set[TopicPartition]): Map[TopicPartition,
Long] = {
+ if (topicOffsets.isEmpty) {
+ partitionOffsets
+ } else {
+ val partitions = assignedPartitions
+ val expanded = partitions.collect {
+ case tp if topicOffsets.contains(tp.topic) => tp ->
topicOffsets(tp.topic)
+ }.toMap
+ val unmatchedTopics =
topicOffsets.keySet.diff(expanded.keySet.map(_.topic))
+ if (unmatchedTopics.nonEmpty) {
+ throw KafkaExceptions.topicOffsetDoesNotMatchAssigned(
+ unmatchedTopics, partitions.map(_.topic))
+ }
+ expanded ++ partitionOffsets
+ }
+ }
+}
/**
* Represents the desire to bind to earliest offset which timestamp for the
offset is equal or
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala
index b8624c5ae637..9233205cbdce 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala
@@ -62,11 +62,14 @@ private[kafka010] trait KafkaOffsetReader {
* This method resolves offset value -1 to the latest and -2 to the
* earliest Kafka seek position.
*
- * @param partitionOffsets the specific offsets to resolve
+ * Offsets bound to a topic as a whole are expanded against the partitions
currently assigned
+ * to the query.
+ *
+ * @param offsets the specific offsets to resolve
* @param reportDataLoss callback to either report or log data loss
depending on setting
*/
def fetchSpecificOffsets(
- partitionOffsets: Map[TopicPartition, Long],
+ offsets: SpecificOffsetRangeLimit,
reportDataLoss: (String, () => Throwable) => Unit): KafkaSourceOffset
/**
@@ -176,6 +179,11 @@ private[kafka010] object KafkaOffsetReader extends Logging
{
private[kafka010] abstract class KafkaOffsetReaderBase extends
KafkaOffsetReader with Logging {
protected val rangeCalculator: KafkaOffsetRangeCalculator
+ /**
+ * @return The set of TopicPartitions currently assigned to the query.
+ */
+ protected def fetchTopicPartitions(): Set[TopicPartition]
+
private def getSortedExecutorList: Array[String] = {
def compare(a: ExecutorCacheTaskLocation, b: ExecutorCacheTaskLocation):
Boolean = {
if (a.host == b.host) {
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 dc22a46e17bc..f6e54c49cc14 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
@@ -124,6 +124,9 @@ private[kafka010] class KafkaOffsetReaderAdmin(
stopAdmin()
}
+ override protected def fetchTopicPartitions(): Set[TopicPartition] =
+ withRetries { resolvePartitions() }
+
override def fetchPartitionOffsets(
offsetRangeLimit: KafkaOffsetRangeLimit,
isStartingOffsets: Boolean): Map[TopicPartition, Long] = {
@@ -135,7 +138,7 @@ private[kafka010] class KafkaOffsetReaderAdmin(
logDebug(s"Assigned partitions: $partitions. Seeking to
$partitionOffsets")
partitionOffsets
}
- val partitions = withRetries { resolvePartitions() }
+ val partitions = fetchTopicPartitions()
// Obtain TopicPartition offsets with late binding support
offsetRangeLimit match {
case EarliestOffsetRangeLimit => partitions.map {
@@ -144,8 +147,8 @@ private[kafka010] class KafkaOffsetReaderAdmin(
case LatestOffsetRangeLimit => partitions.map {
case tp => tp -> KafkaOffsetRangeLimit.LATEST
}.toMap
- case SpecificOffsetRangeLimit(partitionOffsets) =>
- validateTopicPartitions(partitions, partitionOffsets)
+ case offsets: SpecificOffsetRangeLimit =>
+ validateTopicPartitions(partitions, offsets.resolve(partitions))
case SpecificTimestampRangeLimit(partitionTimestamps,
strategyOnNoMatchingStartingOffset) =>
fetchSpecificTimestampBasedOffsets(partitionTimestamps,
isStartingOffsets,
strategyOnNoMatchingStartingOffset).partitionToOffsets
@@ -156,9 +159,12 @@ private[kafka010] class KafkaOffsetReaderAdmin(
}
override def fetchSpecificOffsets(
- partitionOffsets: Map[TopicPartition, Long],
+ offsets: SpecificOffsetRangeLimit,
reportDataLoss: (String, () => Throwable) => Unit): KafkaSourceOffset = {
+ // Topic-level offsets are expanded against the partitions the fetch is
actually run with,
+ // so that metadata changing in between cannot make valid offsets fail the
assertion below
val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = {
partitions =>
+ val partitionOffsets = offsets.resolve(partitions.asScala.toSet)
assert(partitions.asScala == partitionOffsets.keySet,
"If startingOffsets contains specific offsets, you must specify all
TopicPartitions.\n" +
"Use -1 for latest, -2 for earliest, if you don't care.\n" +
@@ -166,8 +172,8 @@ private[kafka010] class KafkaOffsetReaderAdmin(
logDebug(s"Assigned partitions: $partitions. Seeking to
$partitionOffsets")
}
- val fnRetrievePartitionOffsets: ju.Set[TopicPartition] =>
Map[TopicPartition, Long] = { _ =>
- partitionOffsets
+ val fnRetrievePartitionOffsets: ju.Set[TopicPartition] =>
Map[TopicPartition, Long] = {
+ partitions => offsets.resolve(partitions.asScala.toSet)
}
fetchSpecificOffsets0(fnAssertParametersWithPartitions,
fnRetrievePartitionOffsets)
@@ -425,9 +431,11 @@ private[kafka010] class KafkaOffsetReaderAdmin(
// No need to report data loss here
val resolvedFromOffsets =
- fetchSpecificOffsets(fromOffsetsMap, (_, _) => ()).partitionToOffsets
+ fetchSpecificOffsets(SpecificOffsetRangeLimit(fromOffsetsMap), (_, _)
=> ())
+ .partitionToOffsets
val resolvedUntilOffsets =
- fetchSpecificOffsets(untilOffsetsMap, (_, _) => ()).partitionToOffsets
+ fetchSpecificOffsets(SpecificOffsetRangeLimit(untilOffsetsMap), (_, _)
=> ())
+ .partitionToOffsets
val ranges = offsetRangesBase.map(_.topicPartition).map { tp =>
KafkaOffsetRange(tp, resolvedFromOffsets(tp),
resolvedUntilOffsets(tp), preferredLoc = None)
}
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala
index 5196e967399c..3666dfe7541e 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala
@@ -125,10 +125,7 @@ private[kafka010] class KafkaOffsetReaderConsumer(
uninterruptibleThreadRunner.shutdown()
}
- /**
- * @return The Set of TopicPartitions for a given topic
- */
- private def fetchTopicPartitions(): Set[TopicPartition] =
+ override protected def fetchTopicPartitions(): Set[TopicPartition] =
uninterruptibleThreadRunner.runUninterruptibly {
assert(Thread.currentThread().isInstanceOf[UninterruptibleThread])
// Poll to get the latest assigned partitions
@@ -158,8 +155,8 @@ private[kafka010] class KafkaOffsetReaderConsumer(
case LatestOffsetRangeLimit => partitions.map {
case tp => tp -> KafkaOffsetRangeLimit.LATEST
}.toMap
- case SpecificOffsetRangeLimit(partitionOffsets) =>
- validateTopicPartitions(partitions, partitionOffsets)
+ case offsets: SpecificOffsetRangeLimit =>
+ validateTopicPartitions(partitions, offsets.resolve(partitions))
case SpecificTimestampRangeLimit(partitionTimestamps, strategy) =>
fetchSpecificTimestampBasedOffsets(partitionTimestamps,
isStartingOffsets, strategy).partitionToOffsets
@@ -170,9 +167,12 @@ private[kafka010] class KafkaOffsetReaderConsumer(
}
override def fetchSpecificOffsets(
- partitionOffsets: Map[TopicPartition, Long],
+ offsets: SpecificOffsetRangeLimit,
reportDataLoss: (String, () => Throwable) => Unit): KafkaSourceOffset = {
+ // Topic-level offsets are expanded against the partitions the fetch is
actually run with,
+ // so that metadata changing in between cannot make valid offsets fail the
assertion below
val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = {
partitions =>
+ val partitionOffsets = offsets.resolve(partitions.asScala.toSet)
assert(partitions.asScala == partitionOffsets.keySet,
"If startingOffsets contains specific offsets, you must specify all
TopicPartitions.\n" +
"Use -1 for latest, -2 for earliest, if you don't care.\n" +
@@ -180,12 +180,14 @@ private[kafka010] class KafkaOffsetReaderConsumer(
logDebug(s"Partitions assigned to consumer: $partitions. Seeking to
$partitionOffsets")
}
- val fnRetrievePartitionOffsets: ju.Set[TopicPartition] =>
Map[TopicPartition, Long] = { _ =>
- partitionOffsets
+ val fnRetrievePartitionOffsets: ju.Set[TopicPartition] =>
Map[TopicPartition, Long] = {
+ partitions => offsets.resolve(partitions.asScala.toSet)
}
val fnAssertFetchedOffsets: Map[TopicPartition, Long] => Unit = { fetched
=>
- partitionOffsets.foreach {
+ // Only the explicitly specified offsets can be checked here.
Topic-level ones always
+ // expand to the earliest/latest sentinels, which have no expected value
to compare to.
+ offsets.partitionOffsets.foreach {
case (tp, off) if off != KafkaOffsetRangeLimit.LATEST &&
off != KafkaOffsetRangeLimit.EARLIEST =>
if (fetched(tp) != off) {
@@ -466,9 +468,11 @@ private[kafka010] class KafkaOffsetReaderConsumer(
// No need to report data loss here
val resolvedFromOffsets =
- fetchSpecificOffsets(fromOffsetsMap, (_, _) => ()).partitionToOffsets
+ fetchSpecificOffsets(SpecificOffsetRangeLimit(fromOffsetsMap), (_, _)
=> ())
+ .partitionToOffsets
val resolvedUntilOffsets =
- fetchSpecificOffsets(untilOffsetsMap, (_, _) => ()).partitionToOffsets
+ fetchSpecificOffsets(SpecificOffsetRangeLimit(untilOffsetsMap), (_, _)
=> ())
+ .partitionToOffsets
val ranges = offsetRangesBase.map(_.topicPartition).map { tp =>
KafkaOffsetRange(tp, resolvedFromOffsets(tp),
resolvedUntilOffsets(tp), preferredLoc = None)
}
diff --git
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
index 62cd20448453..3ff43cef1fb5 100644
---
a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
+++
b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
@@ -125,7 +125,7 @@ private[kafka010] class KafkaSource(
val offsets = startingOffsets match {
case EarliestOffsetRangeLimit =>
KafkaSourceOffset(kafkaReader.fetchEarliestOffsets())
case LatestOffsetRangeLimit =>
KafkaSourceOffset(kafkaReader.fetchLatestOffsets(None))
- case SpecificOffsetRangeLimit(p) =>
kafkaReader.fetchSpecificOffsets(p, reportDataLoss)
+ case o: SpecificOffsetRangeLimit =>
kafkaReader.fetchSpecificOffsets(o, reportDataLoss)
case SpecificTimestampRangeLimit(p, strategy) =>
kafkaReader.fetchSpecificTimestampBasedOffsets(p, isStartingOffsets
= true, strategy)
case GlobalTimestampRangeLimit(ts, strategy) =>
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 b5a7bcb98fba..e950cdefc35d 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
@@ -375,10 +375,10 @@ private[kafka010] class KafkaSourceProvider extends
DataSourceRegister
case LatestOffsetRangeLimit =>
throw new IllegalArgumentException("starting offset can't be latest " +
"for batch queries on Kafka")
- case SpecificOffsetRangeLimit(partitionOffsets) =>
- partitionOffsets.foreach {
- case (tp, off) if off == KafkaOffsetRangeLimit.LATEST =>
- throw new IllegalArgumentException(s"startingOffsets for $tp can't
" +
+ case SpecificOffsetRangeLimit(partitionOffsets, topicOffsets) =>
+ (partitionOffsets.map { case (tp, off) => tp.toString -> off } ++
topicOffsets).foreach {
+ case (name, off) if off == KafkaOffsetRangeLimit.LATEST =>
+ throw new IllegalArgumentException(s"startingOffsets for $name
can't " +
"be latest for batch queries on Kafka")
case _ => // ignore
}
@@ -393,10 +393,10 @@ private[kafka010] class KafkaSourceProvider extends
DataSourceRegister
throw new IllegalArgumentException("ending offset can't be earliest " +
"for batch queries on Kafka")
case LatestOffsetRangeLimit => // good to go
- case SpecificOffsetRangeLimit(partitionOffsets) =>
- partitionOffsets.foreach {
- case (tp, off) if off == KafkaOffsetRangeLimit.EARLIEST =>
- throw new IllegalArgumentException(s"ending offset for $tp can't
be " +
+ case SpecificOffsetRangeLimit(partitionOffsets, topicOffsets) =>
+ (partitionOffsets.map { case (tp, off) => tp.toString -> off } ++
topicOffsets).foreach {
+ case (name, off) if off == KafkaOffsetRangeLimit.EARLIEST =>
+ throw new IllegalArgumentException(s"ending offset for $name can't
be " +
"earliest for batch queries on Kafka")
case _ => // ignore
}
@@ -640,19 +640,25 @@ private[kafka010] object KafkaSourceProvider extends
Logging {
startOffset match {
case start: SpecificOffsetRangeLimit if
endOffset.isInstanceOf[SpecificOffsetRangeLimit] =>
val end = endOffset.asInstanceOf[SpecificOffsetRangeLimit]
- if (start.partitionOffsets.keySet != end.partitionOffsets.keySet) {
+ // Topic-level offsets are only expanded once the partitions are
discovered, so with them
+ // the two sides can legitimately enumerate different topic-partitions
here. Matching them
+ // up is then left to the offset readers, which see the assigned
partitions.
+ val enumeratesAllPartitions = start.topicOffsets.isEmpty &&
end.topicOffsets.isEmpty
+ if (enumeratesAllPartitions &&
+ start.partitionOffsets.keySet != end.partitionOffsets.keySet) {
throw KafkaExceptions.unmatchedTopicPartitionsBetweenOffsets(
start.partitionOffsets.keySet, end.partitionOffsets.keySet
)
}
start.partitionOffsets.foreach {
- case (tp, startOffset) =>
+ case (tp, startOffset) if end.partitionOffsets.contains(tp) =>
checkStartOffsetNotGreaterThanEndOffset(
startOffset,
end.partitionOffsets(tp),
tp,
KafkaExceptions.unresolvedStartOffsetGreaterThanEndOffset
)
+ case _ => // the offsets of this partition are not comparable before
resolution
}
case start: SpecificTimestampRangeLimit
@@ -711,7 +717,7 @@ private[kafka010] object KafkaSourceProvider extends
Logging {
LatestOffsetRangeLimit
case Some(offset) if offset.toLowerCase(Locale.ROOT) == "earliest" =>
EarliestOffsetRangeLimit
- case Some(json) =>
SpecificOffsetRangeLimit(JsonUtils.partitionOffsets(json))
+ case Some(json) => JsonUtils.specificOffsets(json)
case None => defaultOffsets
}
}
diff --git
a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala
index 54b980049d1a..e7b08c257556 100644
---
a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala
+++
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala
@@ -42,4 +42,59 @@ class JsonUtilsSuite extends SparkFunSuite {
assert(parsed(new TopicPartition("topicA", 1)) === -1)
assert(parsed(new TopicPartition("topicB", 0)) === -2)
}
+
+ test("parsing specificOffsets") {
+ val parsed = JsonUtils.specificOffsets(
+ """{"topicA":{"0":23,"1":-1},"topicB":{"0":-2}}""")
+
+ assert(parsed.partitionOffsets === Map(
+ new TopicPartition("topicA", 0) -> 23L,
+ new TopicPartition("topicA", 1) -> -1L,
+ new TopicPartition("topicB", 0) -> -2L))
+ assert(parsed.topicOffsets.isEmpty)
+ }
+
+ test("parsing specificOffsets with topic-level earliest/latest") {
+ val parsed = JsonUtils.specificOffsets(
+ """{"topicA":"earliest","topicB":"latest"}""")
+
+ assert(parsed.partitionOffsets.isEmpty)
+ assert(parsed.topicOffsets === Map(
+ "topicA" -> KafkaOffsetRangeLimit.EARLIEST,
+ "topicB" -> KafkaOffsetRangeLimit.LATEST))
+ }
+
+ test("parsing specificOffsets mixing topic-level and partition-level
values") {
+ val parsed = JsonUtils.specificOffsets(
+ """{"topicA":"earliest","topicB":{"0":23,"1":-1},"topicC":"LATEST"}""")
+
+ assert(parsed.partitionOffsets === Map(
+ new TopicPartition("topicB", 0) -> 23L,
+ new TopicPartition("topicB", 1) -> -1L))
+ assert(parsed.topicOffsets === Map(
+ "topicA" -> KafkaOffsetRangeLimit.EARLIEST,
+ "topicC" -> KafkaOffsetRangeLimit.LATEST))
+ }
+
+ test("parsing specificOffsets rejects a topic given in both forms") {
+ val ex = intercept[IllegalArgumentException] {
+ JsonUtils.specificOffsets("""{"topicA":"earliest","topicA":{"0":23}}""")
+ }
+ assert(ex.getMessage.contains("topicA"))
+ assert(ex.getMessage.contains("both a topic-level offset"))
+ }
+
+ test("parsing specificOffsets rejects malformed json") {
+ Seq(
+ """{"topicA":"first"}""",
+ """{"topicA":[0,1]}""",
+ """{"topicA":{"0":"earliest"}}""",
+ """"earliest"""",
+ "not json").foreach { json =>
+ val ex = intercept[IllegalArgumentException] {
+ JsonUtils.specificOffsets(json)
+ }
+ assert(ex.getMessage.contains(json))
+ }
+ }
}
diff --git
a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala
index 274be623da6f..66a1bf00f87f 100644
---
a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala
+++
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala
@@ -911,6 +911,54 @@ abstract class KafkaMicroBatchSourceSuiteBase extends
KafkaSourceSuiteBase with
)
}
+ test("topic-level offsets tolerate a new topic matching subscribePattern
after the start") {
+ val topicPrefix = newTopic()
+ val topic = s"$topicPrefix-a"
+ val topic2 = s"$topicPrefix-b"
+ testUtils.createTopic(topic, partitions = 1)
+ testUtils.sendMessages(topic, Array("1"), Some(0))
+
+ // The starting offsets only name the topic that exists when the query
starts
+ val ds = spark
+ .readStream
+ .format("kafka")
+ .option("kafka.bootstrap.servers", testUtils.brokerAddress)
+ .option("kafka.metadata.max.age.ms", "1")
+ .option("subscribePattern", s"$topicPrefix-.*")
+ .option("startingOffsets", s"""{"$topic":"earliest"}""")
+ .load()
+ .selectExpr("CAST(value AS STRING)")
+ .as[String]
+ .map(_.toInt)
+
+ testStream(ds)(
+ StartStream(),
+ AssertOnQuery { q =>
+ q.processAllAvailable()
+ true
+ },
+ CheckAnswer(1),
+ // A topic created after the initial offsets were resolved is picked up
as a new partition,
+ // starting at earliest, instead of tripping the strict topic check
+ WithOffsetSync(new TopicPartition(topic2, 0), expectedOffset = 1) { () =>
+ testUtils.createTopic(topic2, partitions = 1)
+ testUtils.sendMessages(topic2, Array("2"), Some(0))
+ },
+ AssertOnQuery { q =>
+ // The consumer based reader only sees a newly created topic on its
next metadata refresh,
+ // so keep triggering batches until the topic shows up in the query's
offsets
+ eventually(timeout(streamingTimeout)) {
+ q.processAllAvailable()
+ val progress = q.lastProgress
+ assert(progress != null &&
progress.sources.exists(_.endOffset.contains(topic2)),
+ s"$topic2 has not been discovered yet")
+ }
+ true
+ },
+ CheckAnswer(1, 2)
+ )
+ }
+
test("ensure that initial offset are written with an extra byte in the
beginning (SPARK-19517)") {
withTempDir { metadataPath =>
val topic = "kafka-initial-offset-current"
@@ -2100,6 +2148,42 @@ abstract class KafkaSourceSuiteBase extends
KafkaSourceTest {
}
}
+ test("subscribing topics from topic-level offsets") {
+ val topic1 = newTopic()
+ val topic2 = newTopic()
+ testUtils.createTopic(topic1, partitions = 3)
+ testUtils.createTopic(topic2, partitions = 2)
+ testUtils.sendMessages(topic1, (1 to 3).map(_.toString).toArray, Some(0))
+ testUtils.sendMessages(topic2, Array("11"), Some(0))
+
+ // topic1 is read from its beginning while topic2 only contributes the
records added after
+ // the query started, without either topic having its partitions
enumerated in the option.
+ val kafka = spark
+ .readStream
+ .format("kafka")
+ .option("kafka.bootstrap.servers", testUtils.brokerAddress)
+ .option("kafka.metadata.max.age.ms", "1")
+ .option("subscribe", s"$topic1,$topic2")
+ .option("startingOffsets",
s"""{"$topic1":"earliest","$topic2":"latest"}""")
+ .load()
+ .selectExpr("CAST(value AS STRING)")
+ .as[String]
+ val mapped = kafka.map(_.toInt)
+
+ testStream(mapped)(
+ makeSureGetOffsetCalled,
+ // Records written to topic2 after the query started are read, unlike
"11" which predates it
+ WithOffsetSync(new TopicPartition(topic2, 0), expectedOffset = 2) { () =>
+ testUtils.sendMessages(topic2, Array("12"), Some(0))
+ },
+ AddKafkaData(Set(topic1, topic2), 4),
+ CheckAnswer(1, 2, 3, 12, 4),
+ StopStream,
+ StartStream(),
+ CheckAnswer(1, 2, 3, 12, 4) // Should get the data back on recovery
+ )
+ }
+
test("subscribing topic by name from specific timestamps with non-matching
starting offset") {
val topic = newTopic()
testFromSpecificTimestampsWithNoMatchingStartingOffset(topic, "subscribe"
-> topic)
@@ -2339,7 +2423,11 @@ abstract class KafkaSourceSuiteBase extends
KafkaSourceTest {
(STARTING_OFFSETS_OPTION_KEY, "earLiEst", EarliestOffsetRangeLimit),
(ENDING_OFFSETS_OPTION_KEY, "laTest", LatestOffsetRangeLimit),
(STARTING_OFFSETS_OPTION_KEY, """{"topic-A":{"0":23}}""",
- SpecificOffsetRangeLimit(Map(new TopicPartition("topic-A", 0) ->
23))))) {
+ SpecificOffsetRangeLimit(Map(new TopicPartition("topic-A", 0) -> 23))),
+ (STARTING_OFFSETS_OPTION_KEY,
"""{"topic-A":"eArLiEst","topic-B":{"0":23}}""",
+ SpecificOffsetRangeLimit(
+ Map(new TopicPartition("topic-B", 0) -> 23),
+ Map("topic-A" -> KafkaOffsetRangeLimit.EARLIEST))))) {
val offset = getKafkaOffsetRangeLimit(
CaseInsensitiveMap[String](Map(optionKey -> optionValue)), "dummy",
"dummy", optionKey,
answer)
diff --git
a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimitSuite.scala
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimitSuite.scala
new file mode 100644
index 000000000000..8041d7e56d92
--- /dev/null
+++
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimitSuite.scala
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.kafka010
+
+import org.apache.kafka.common.TopicPartition
+
+import org.apache.spark.SparkFunSuite
+
+class KafkaOffsetRangeLimitSuite extends SparkFunSuite {
+
+ test("resolving topic-level offsets against the assigned partitions") {
+ val limit = SpecificOffsetRangeLimit(
+ Map(new TopicPartition("topicB", 0) -> 23L),
+ Map("topicA" -> KafkaOffsetRangeLimit.EARLIEST))
+ val partitions = Set(
+ new TopicPartition("topicA", 0),
+ new TopicPartition("topicA", 1),
+ new TopicPartition("topicB", 0))
+
+ assert(limit.resolve(partitions) === Map(
+ new TopicPartition("topicA", 0) -> KafkaOffsetRangeLimit.EARLIEST,
+ new TopicPartition("topicA", 1) -> KafkaOffsetRangeLimit.EARLIEST,
+ new TopicPartition("topicB", 0) -> 23L))
+ }
+
+ test("resolving fully enumerated offsets doesn't need the assigned
partitions") {
+ val limit = SpecificOffsetRangeLimit(Map(new TopicPartition("topicB", 0)
-> 23L))
+ assert(limit.resolve(fail("the partitions should not have been fetched"))
===
+ Map(new TopicPartition("topicB", 0) -> 23L))
+ }
+
+ test("resolving a topic-level offset of a topic without assigned partitions
fails") {
+ val limit = SpecificOffsetRangeLimit(
+ Map.empty,
+ Map("topicA" -> KafkaOffsetRangeLimit.EARLIEST,
+ "topicC" -> KafkaOffsetRangeLimit.EARLIEST))
+ val ex = intercept[KafkaIllegalStateException] {
+ limit.resolve(Set(new TopicPartition("topicA", 0)))
+ }
+ assert(ex.getCondition === "KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED")
+ assert(ex.getMessage.contains("topicC"))
+ }
+}
diff --git
a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala
index ab515e2a8e49..3b71abf026d2 100644
---
a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala
+++
b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala
@@ -25,7 +25,7 @@ import org.apache.kafka.clients.producer.ProducerRecord
import org.apache.kafka.common.TopicPartition
import org.apache.spark.{SparkConf, TestUtils}
-import org.apache.spark.sql.DataFrameReader
+import org.apache.spark.sql.{DataFrame, DataFrameReader}
import org.apache.spark.sql.execution.datasources.LogicalRelation
import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
import org.apache.spark.sql.internal.SQLConf
@@ -151,6 +151,105 @@ abstract class KafkaRelationSuiteBase extends
SharedSparkSession with KafkaTest
checkAnswer(df, (0 to 30).map(_.toString).toDF())
}
+ test("topic-level offsets") {
+ val topic1 = newTopic()
+ val topic2 = newTopic()
+ testUtils.createTopic(topic1, partitions = 3)
+ testUtils.createTopic(topic2, partitions = 2)
+ testUtils.sendMessages(topic1, (0 to 9).map(_.toString).toArray, Some(0))
+ testUtils.sendMessages(topic1, (10 to 19).map(_.toString).toArray, Some(1))
+ testUtils.sendMessages(topic1, Array("20"), Some(2))
+ testUtils.sendMessages(topic2, (100 to 109).map(_.toString).toArray,
Some(0))
+ testUtils.sendMessages(topic2, (110 to 119).map(_.toString).toArray,
Some(1))
+
+ // A topic bound as a whole covers all of its partitions, without
enumerating them. This is a
+ // def so that every query plans against the partitions discovered at that
point.
+ def df: DataFrame = createDF(s"$topic1,$topic2", withOptions = Map(
+ "kafka.metadata.max.age.ms" -> "1",
+ "startingOffsets" -> s"""{"$topic1":"earliest","$topic2":"earliest"}""",
+ "endingOffsets" -> s"""{"$topic1":"latest","$topic2":"latest"}"""))
+ checkAnswer(df, ((0 to 20) ++ (100 to 119)).map(_.toString).toDF())
+
+ // Topic-level offsets are expanded against the partitions discovered when
the query is
+ // planned, so a partition added afterwards is picked up without touching
the options
+ testUtils.addPartitions(topic2, 3)
+ testUtils.sendMessages(topic2, Array("120"), Some(2))
+ checkAnswer(df, ((0 to 20) ++ (100 to 120)).map(_.toString).toDF())
+ }
+
+ test("mixed topic-level and partition-level offsets") {
+ val topic1 = newTopic()
+ val topic2 = newTopic()
+ testUtils.createTopic(topic1, partitions = 2)
+ testUtils.createTopic(topic2, partitions = 2)
+ testUtils.sendMessages(topic1, (0 to 9).map(_.toString).toArray, Some(0))
+ testUtils.sendMessages(topic1, (10 to 19).map(_.toString).toArray, Some(1))
+ testUtils.sendMessages(topic2, (100 to 109).map(_.toString).toArray,
Some(0))
+ testUtils.sendMessages(topic2, (110 to 119).map(_.toString).toArray,
Some(1))
+
+ // topic1 starts from its earliest offsets, while topic2 is enumerated per
partition
+ val df = createDF(s"$topic1,$topic2", withOptions = Map(
+ "startingOffsets" ->
s"""{"$topic1":"earliest","$topic2":{"0":5,"1":-2}}""",
+ "endingOffsets" -> s"""{"$topic1":"latest","$topic2":{"0":-1,"1":8}}"""))
+ checkAnswer(df, ((0 to 19) ++ (105 to 109) ++ (110 to
117)).map(_.toString).toDF())
+ }
+
+ test("topic-level offsets with subscribePattern") {
+ val prefix = newTopic()
+ val topic1 = s"$prefix-a"
+ val topic2 = s"$prefix-b"
+ testUtils.createTopic(topic1, partitions = 2)
+ testUtils.createTopic(topic2, partitions = 2)
+ testUtils.sendMessages(topic1, (0 to 9).map(_.toString).toArray, Some(0))
+ testUtils.sendMessages(topic1, (10 to 19).map(_.toString).toArray, Some(1))
+ testUtils.sendMessages(topic2, (100 to 109).map(_.toString).toArray,
Some(0))
+ testUtils.sendMessages(topic2, (110 to 119).map(_.toString).toArray,
Some(1))
+
+ val df = spark
+ .read
+ .format("kafka")
+ .option("kafka.bootstrap.servers", testUtils.brokerAddress)
+ .option("subscribePattern", s"$prefix-.*")
+ .option("startingOffsets",
s"""{"$topic1":"earliest","$topic2":{"0":5,"1":-2}}""")
+ .option("endingOffsets",
s"""{"$topic1":"latest","$topic2":{"0":-1,"1":8}}""")
+ .load()
+ .selectExpr("CAST(value AS STRING)")
+ checkAnswer(df, ((0 to 19) ++ (105 to 109) ++ (110 to
117)).map(_.toString).toDF())
+ }
+
+ test("topic-level offsets must line up with the topics matched by
subscribePattern") {
+ val prefix = newTopic()
+ val topic1 = s"$prefix-a"
+ val topic2 = s"$prefix-b"
+ testUtils.createTopic(topic1, partitions = 1)
+ testUtils.createTopic(topic2, partitions = 1)
+ testUtils.sendMessages(topic1, Array("1"), Some(0))
+ testUtils.sendMessages(topic2, Array("2"), Some(0))
+
+ def readWith(startingOffsets: String): Unit = {
+ spark
+ .read
+ .format("kafka")
+ .option("kafka.bootstrap.servers", testUtils.brokerAddress)
+ .option("subscribePattern", s"$prefix-.*")
+ .option("startingOffsets", startingOffsets)
+ .load()
+ .collect()
+ }
+
+ // A topic matched by the pattern must still be covered by the offsets
+ val uncovered = intercept[KafkaIllegalStateException] {
+ readWith(s"""{"$topic1":"earliest"}""")
+ }
+ assert(uncovered.getCondition ===
"KAFKA_START_OFFSET_DOES_NOT_MATCH_ASSIGNED")
+
+ // ... and a topic-level offset for a topic the pattern doesn't match is
rejected
+ val unknown = intercept[KafkaIllegalStateException] {
+
readWith(s"""{"$topic1":"earliest","$topic2":"earliest","$prefix-ghost":"earliest"}""")
+ }
+ assert(unknown.getCondition ===
"KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED")
+ }
+
test("default starting and ending offsets with headers") {
val topic = newTopic()
testUtils.createTopic(topic, partitions = 3)
@@ -470,6 +569,9 @@ abstract class KafkaRelationSuiteBase extends
SharedSparkSession with KafkaTest
testBadOptions("subscribe" -> "t", "startingOffsets" -> startingOffsets)(
"startingOffsets for t-0 can't be latest for batch queries on Kafka")
+ // Now do it with a topic-level start offset indicating latest
+ testBadOptions("subscribe" -> "t", "startingOffsets" ->
"""{"t":"latest"}""")(
+ "startingOffsets for t can't be latest for batch queries on Kafka")
// Make sure we catch ending offsets that indicate earliest
testBadOptions("endingOffsets" -> "earliest")("ending offset can't be
earliest " +
@@ -481,6 +583,10 @@ abstract class KafkaRelationSuiteBase extends
SharedSparkSession with KafkaTest
testBadOptions("subscribe" -> "t", "endingOffsets" -> endingOffsets)(
"ending offset for t-0 can't be earliest for batch queries on Kafka")
+ // Make sure we catch a topic-level ending offset indicating earliest
+ testBadOptions("subscribe" -> "t", "endingOffsets" ->
"""{"t":"earliest"}""")(
+ "ending offset for t can't be earliest for batch queries on Kafka")
+
// No strategy specified
testBadOptions()("options must be specified", "subscribe",
"subscribePattern")
diff --git a/docs/streaming/structured-streaming-kafka-integration.md
b/docs/streaming/structured-streaming-kafka-integration.md
index 105dac7fb4fc..5be6dfa74173 100644
--- a/docs/streaming/structured-streaming-kafka-integration.md
+++ b/docs/streaming/structured-streaming-kafka-integration.md
@@ -403,14 +403,24 @@ The following configurations are optional:
<tr>
<td>startingOffsets</td>
<td>"earliest", "latest" (streaming only), or json string
- """ {"topicA":{"0":23,"1":-1},"topicB":{"0":-2}} """
+ """ {"topicA":{"0":23,"1":-1},"topicB":{"0":-2}} """ or
+ """ {"topicA":"earliest","topicB":"latest"} """ or a mix of both forms
+ """ {"topicA":"earliest","topicB":{"0":23,"1":-1},"topicC":"latest"} """
</td>
<td>"latest" for streaming, "earliest" for batch</td>
<td>streaming and batch</td>
<td>The start point when a query is started, either "earliest" which is from
the earliest offsets,
"latest" which is just from the latest offsets, or a json string specifying
a starting offset for
each TopicPartition. In the json, -2 as an offset can be used to refer to
earliest, -1 to latest.
- Note: For batch queries, latest (either implicitly or by using -1 in json)
is not allowed.
+ A topic may also map to the string "earliest" or "latest" instead of an
object, which applies that
+ offset to every partition of the topic without having to enumerate them; the
two forms can be mixed
+ in the same json. Topic-level values are expanded against the partitions
discovered when the offsets
+ are resolved, so they keep working when a topic is repartitioned. As with
the per-partition form,
+ the json must account for every topic subscribed at the time the offsets are
resolved, so with
+ <code>subscribePattern</code> the topics matched by the pattern have to be
known then; topics
+ matched later on are discovered as new partitions and start at earliest, as
usual.
+ Note: For batch queries, latest (either implicitly, by using -1 in json, or
by binding a topic to
+ "latest") is not allowed.
For streaming queries, this only applies when a new query is started, and
that resuming will
always pick up from where the query left off. Newly discovered partitions
during a query will start at
earliest.</td>
@@ -442,13 +452,16 @@ The following configurations are optional:
<tr>
<td>endingOffsets</td>
<td>latest or json string
- {"topicA":{"0":23,"1":-1},"topicB":{"0":-1}}
+ {"topicA":{"0":23,"1":-1},"topicB":{"0":-1}} or
{"topicA":"latest","topicB":"latest"}
+ or a mix of both forms {"topicA":"latest","topicB":{"0":23,"1":-1}}
</td>
<td>latest</td>
<td>batch query</td>
<td>The end point when a batch query is ended, either "latest" which is just
referred to the
latest, or a json string specifying an ending offset for each
TopicPartition. In the json, -1
- as an offset can be used to refer to latest, and -2 (earliest) as an offset
is not allowed.</td>
+ as an offset can be used to refer to latest, and -2 (earliest) as an offset
is not allowed.
+ As with <code>startingOffsets</code>, a topic may map to the string "latest"
instead of an object to
+ apply that offset to every partition of the topic; binding a topic to
"earliest" is not allowed.</td>
</tr>
<tr>
<td>failOnDataLoss</td>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
