[SPARK-29223][SQL][SS] New option to specify timestamp on all subscribing topic-partitions in Kafka source

### What changes were proposed in this pull request?

This patch is a follow-up of SPARK-26848 (#23747). In SPARK-26848, we decided to open possibility to let end users set individual timestamp per partition. But in many cases, specifying timestamp represents the intention that we would want to go back to specific timestamp and reprocess records, which should be applied to all topics and partitions.

This patch proposes to provide a way to set a global timestamp across topic-partitions which the source is subscribing to, so that end users can set all offsets by specific timestamp easily. To provide the way to config the timestamp easier, the new options only receive "a" timestamp for start/end timestamp.

New options introduced in this PR:

* startingTimestamp
* endingTimestamp

All two options receive timestamp as string.

There're priorities for options regarding starting/ending offset as we will have three options for start offsets and another three options for end offsets. Priorities are following:

* starting offsets: startingTimestamp -> startingOffsetsByTimestamp -> startingOffsets
* ending offsets: startingTimestamp -> startingOffsetsByTimestamp -> startingOffsets

### Why are the changes needed?

Existing option to specify timestamp as offset is quite verbose if there're a lot of partitions across topics. Suppose there're 100s of partitions in a topic, the json should contain 100s of times of the same timestamp.

Also, the number of partitions can also change, which requires either:

* fixing the code if the json is statically created
* introducing the dependencies on Kafka client and deal with Kafka API on crafting json programmatically

Both approaches are even not "acceptable" if we're dealing with ad-hoc query; anyone doesn't want to write the code more complicated than the query itself. Flink [provides the option](https://ci.apache.org/projects/flink/flink-docs-release-1.13/docs/connectors/datastream/kafka/#kafka-consumers-start-position-configuration) to specify a timestamp for all topic-partitions like this PR, and even doesn't provide the option to specify the timestamp per topic-partition.

With this PR, end users are only required to provide a single timestamp value. No more complicated JSON format end users need to know about the structure.

### Does this PR introduce _any_ user-facing change?

Yes, this PR introduces two new options, described in above section.

Doc changes are following:

![스크린샷 2021-05-21 오후 12 01 02](https://user-images.githubusercontent.com/1317309/119076244-3034e680-ba2d-11eb-8323-0e227932d2e5.png)
![스크린샷 2021-05-21 오후 12 01 12](https://user-images.githubusercontent.com/1317309/119076255-35923100-ba2d-11eb-9d79-538a7f9ee738.png)
![스크린샷 2021-05-21 오후 12 01 24](https://user-images.githubusercontent.com/1317309/119076264-39be4e80-ba2d-11eb-8265-ac158f55c360.png)
![스크린샷 2021-05-21 오후 12 06 01](https://user-images.githubusercontent.com/1317309/119076271-3d51d580-ba2d-11eb-98ea-35fd72b1bbfc.png)

### How was this patch tested?

New UTs covering new functionalities. Also manually tested via simple batch & streaming queries.

Closes #32609 from HeartSaVioR/SPARK-29223-v2.

Authored-by: Jungtaek Lim <kabhwan.opensource@gmail.com>
Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

Author: HeartSaVioR

docs/structured-streaming-kafka-integration.md

@@ -362,24 +362,33 @@ The following configurations are optional:
 
 <table class="table">
 <tr><th>Option</th><th>value</th><th>default</th><th>query type</th><th>meaning</th></tr>
+<tr>
+  <td>startingTimestamp</td>
+  <td>timestamp string e.g. "1000"</td>
+  <td>none (next preference is <code>startingOffsetsByTimestamp</code>)</td>
+  <td>streaming and batch</td>
+  <td>The start point of timestamp when a query is started, a string specifying a starting timestamp for
+  all partitions in topics being subscribed. Please refer the details on timestamp offset options below. If Kafka doesn't return the matched offset,
+  the query will fail immediately to prevent unintended read from such partition. (This is a kind of limitation as of now, and will be addressed in near future.)<p/>
+  <p/>
+  Note1: <code>startingTimestamp</code> takes precedence over <code>startingOffsetsByTimestamp</code> and <code>startingOffsets</code>.<p/>
+  Note2: 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>
+</tr>
 <tr>
   <td>startingOffsetsByTimestamp</td>
   <td>json string
   """ {"topicA":{"0": 1000, "1": 1000}, "topicB": {"0": 2000, "1": 2000}} """
   </td>
-  <td>none (the value of <code>startingOffsets</code> will apply)</td>
+  <td>none (next preference is <code>startingOffsets</code>)</td>
   <td>streaming and batch</td>
   <td>The start point of timestamp when a query is started, a json string specifying a starting timestamp for
-  each TopicPartition. The returned offset for each partition is the earliest offset whose timestamp is greater than or
-  equal to the given timestamp in the corresponding partition. If the matched offset doesn't exist,
+  each TopicPartition. Please refer the details on timestamp offset options below. If Kafka doesn't return the matched offset,
   the query will fail immediately to prevent unintended read from such partition. (This is a kind of limitation as of now, and will be addressed in near future.)<p/>
   <p/>
-  Spark simply passes the timestamp information to <code>KafkaConsumer.offsetsForTimes</code>, and doesn't interpret or reason about the value. <p/>
-  For more details on <code>KafkaConsumer.offsetsForTimes</code>, please refer <a href="https://kafka.apache.org/21/javadoc/org/apache/kafka/clients/consumer/KafkaConsumer.html#offsetsForTimes-java.util.Map-">javadoc</a> for details.<p/>
-  Also the meaning of <code>timestamp</code> here can be vary according to Kafka configuration (<code>log.message.timestamp.type</code>): please refer <a href="https://kafka.apache.org/documentation/">Kafka documentation</a> for further details.<p/>
-  Note: This option requires Kafka 0.10.1.0 or higher.<p/>
-  Note2: <code>startingOffsetsByTimestamp</code> takes precedence over <code>startingOffsets</code>.<p/>
-  Note3: For streaming queries, this only applies when a new query is started, and that resuming will
+  Note1: <code>startingOffsetsByTimestamp</code> takes precedence over <code>startingOffsets</code>.<p/>
+  Note2: 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>
 </tr>
@@ -398,23 +407,28 @@ The following configurations are optional:
   always pick up from where the query left off. Newly discovered partitions during a query will start at
   earliest.</td>
 </tr>
+<tr>
+  <td>endingTimestamp</td>
+  <td>timestamp string e.g. "1000"</td>
+  <td>none (next preference is <code>endingOffsetsByTimestamp</code>)</td>
+  <td>batch query</td>
+  <td>The end point when a batch query is ended, a json string specifying an ending timestamp for
+  all partitions in topics being subscribed. Please refer the details on timestamp offset options below.
+  If Kafka doesn't return the matched offset, the offset will be set to latest.<p/>
+  Note: <code>endingTimestamp</code> takes precedence over <code>endingOffsetsByTimestamp</code> and <code>endingOffsets</code>.<p/>
+  </td>
+</tr>
 <tr>
   <td>endingOffsetsByTimestamp</td>
   <td>json string
   """ {"topicA":{"0": 1000, "1": 1000}, "topicB": {"0": 2000, "1": 2000}} """
   </td>
-  <td>latest</td>
+  <td>none (next preference is <code>endingOffsets</code>)</td>
   <td>batch query</td>
   <td>The end point when a batch query is ended, a json string specifying an ending timestamp for each TopicPartition.
-  The returned offset for each partition is the earliest offset whose timestamp is greater than or equal to
-  the given timestamp in the corresponding partition. If the matched offset doesn't exist, the offset will
-  be set to latest.<p/>
-  <p/>
-  Spark simply passes the timestamp information to <code>KafkaConsumer.offsetsForTimes</code>, and doesn't interpret or reason about the value. <p/>
-  For more details on <code>KafkaConsumer.offsetsForTimes</code>, please refer <a href="https://kafka.apache.org/21/javadoc/org/apache/kafka/clients/consumer/KafkaConsumer.html#offsetsForTimes-java.util.Map-">javadoc</a> for details.<p/>
-  Also the meaning of <code>timestamp</code> here can be vary according to Kafka configuration (<code>log.message.timestamp.type</code>): please refer <a href="https://kafka.apache.org/documentation/">Kafka documentation</a> for further details.<p/>
-  Note: This option requires Kafka 0.10.1.0 or higher.<p/>
-  Note2: <code>endingOffsetsByTimestamp</code> takes precedence over <code>endingOffsets</code>.
+  Please refer the details on timestamp offset options below. If Kafka doesn't return the matched offset,
+  the offset will be set to latest.<p/>
+  Note: <code>endingOffsetsByTimestamp</code> takes precedence over <code>endingOffsets</code>.
   </td>
 </tr>
 <tr>
@@ -512,6 +526,17 @@ The following configurations are optional:
 </tr>
 </table>
 
+### Details on timestamp offset options
+
+The returned offset for each partition is the earliest offset whose timestamp is greater than or equal to the given timestamp in the corresponding partition.
+The behavior varies across options if Kafka doesn't return the matched offset - check the description of each option.
+
+Spark simply passes the timestamp information to <code>KafkaConsumer.offsetsForTimes</code>, and doesn't interpret or reason about the value.
+For more details on <code>KafkaConsumer.offsetsForTimes</code>, please refer <a href="http://kafka.apache.org/0101/javadoc/org/apache/kafka/clients/consumer/KafkaConsumer.html#offsetsForTimes(java.util.Map)">javadoc</a> for details.
+Also, the meaning of <code>timestamp</code> here can be vary according to Kafka configuration (<code>log.message.timestamp.type</code>): please refer <a href="https://kafka.apache.org/documentation/">Kafka documentation</a> for further details.
+
+Timestamp offset options require Kafka 0.10.1.0 or higher.
+
 ### Offset fetching
 
 In Spark 3.0 and before Spark uses <code>KafkaConsumer</code> for offset fetching which could cause infinite wait in the driver.

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala

@@ -72,6 +72,8 @@ class KafkaContinuousStream(
       case SpecificOffsetRangeLimit(p) => offsetReader.fetchSpecificOffsets(p, reportDataLoss)
       case SpecificTimestampRangeLimit(p) => offsetReader.fetchSpecificTimestampBasedOffsets(p,
         failsOnNoMatchingOffset = true)
+      case GlobalTimestampRangeLimit(ts) => offsetReader.fetchGlobalTimestampBasedOffsets(
+        ts, failsOnNoMatchingOffset = true)
     }
     logInfo(s"Initial offsets: $offsets")
     offsets

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala

@@ -166,6 +166,8 @@ private[kafka010] class KafkaMicroBatchStream(
           kafkaOffsetReader.fetchSpecificOffsets(p, reportDataLoss)
         case SpecificTimestampRangeLimit(p) =>
           kafkaOffsetReader.fetchSpecificTimestampBasedOffsets(p, failsOnNoMatchingOffset = true)
+        case GlobalTimestampRangeLimit(ts) =>
+          kafkaOffsetReader.fetchGlobalTimestampBasedOffsets(ts, failsOnNoMatchingOffset = true)
       }
       metadataLog.add(0, offsets)
       logInfo(s"Initial offsets: $offsets")

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala

@@ -49,6 +49,13 @@ private[kafka010] case class SpecificOffsetRangeLimit(
 private[kafka010] case class SpecificTimestampRangeLimit(
     topicTimestamps: Map[TopicPartition, Long]) extends KafkaOffsetRangeLimit
 
+/**
+ * Represents the desire to bind to earliest offset which timestamp for the offset is equal or
+ * greater than specific timestamp. This applies the timestamp to the all topics/partitions.
+ */
+private[kafka010] case class GlobalTimestampRangeLimit(
+    timestamp: Long) extends KafkaOffsetRangeLimit
+
 private[kafka010] object KafkaOffsetRangeLimit {
   /**
    * Used to denote offset range limits that are resolved via Kafka

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala

@@ -77,6 +77,20 @@ private[kafka010] trait KafkaOffsetReader {
       partitionTimestamps: Map[TopicPartition, Long],
       failsOnNoMatchingOffset: Boolean): KafkaSourceOffset
 
+  /**
+   * Resolves the specific offsets based on timestamp per all topic-partitions being subscribed.
+   * The returned offset for each partition is the earliest offset whose timestamp is greater
+   * than or equal to the given timestamp in the corresponding partition. If the matched offset
+   * doesn't exist, depending on `failsOnNoMatchingOffset` parameter, the offset will be set to
+   * latest or this method throws an error.
+   *
+   * @param timestamp the timestamp.
+   * @param failsOnNoMatchingOffset whether to fail the query when no matched offset can be found.
+   */
+  def fetchGlobalTimestampBasedOffsets(
+      timestamp: Long,
+      failsOnNoMatchingOffset: Boolean): KafkaSourceOffset
+
   /**
    * Fetch the earliest offsets for the topic partitions that are indicated
    * in the [[ConsumerStrategy]].

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

@@ -138,6 +138,9 @@ private[kafka010] class KafkaOffsetReaderAdmin(
       case SpecificTimestampRangeLimit(partitionTimestamps) =>
         fetchSpecificTimestampBasedOffsets(partitionTimestamps,
           failsOnNoMatchingOffset = isStartingOffsets).partitionToOffsets
+      case GlobalTimestampRangeLimit(timestamp) =>
+        fetchGlobalTimestampBasedOffsets(timestamp,
+          failsOnNoMatchingOffset = isStartingOffsets).partitionToOffsets
     }
   }
 
@@ -193,6 +196,37 @@ private[kafka010] class KafkaOffsetReaderAdmin(
     fetchSpecificOffsets0(fnAssertParametersWithPartitions, fnRetrievePartitionOffsets)
   }
 
+
+  override def fetchGlobalTimestampBasedOffsets(
+      timestamp: Long,
+      failsOnNoMatchingOffset: Boolean): KafkaSourceOffset = {
+    val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = { partitions =>
+      logDebug(s"Partitions assigned to consumer: $partitions. Seeking to $timestamp")
+    }
+
+    val fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long] = { tps =>
+      val listOffsetsParams = tps.asScala.map { tp =>
+        tp -> OffsetSpec.forTimestamp(timestamp)
+      }.toMap.asJava
+      admin.listOffsets(listOffsetsParams, listOffsetsOptions).all().get().asScala.map {
+        case (tp, offsetSpec) =>
+          if (failsOnNoMatchingOffset) {
+            assert(offsetSpec.offset() != OffsetFetchResponse.INVALID_OFFSET, "No offset " +
+              s"matched from request of topic-partition $tp and timestamp " +
+              s"$timestamp.")
+          }
+
+          if (offsetSpec.offset() == OffsetFetchResponse.INVALID_OFFSET) {
+            tp -> KafkaOffsetRangeLimit.LATEST
+          } else {
+            tp -> offsetSpec.offset()
+          }
+      }.toMap
+    }
+
+    fetchSpecificOffsets0(fnAssertParametersWithPartitions, fnRetrievePartitionOffsets)
+  }
+
   private def fetchSpecificOffsets0(
       fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit,
       fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long]

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala

@@ -160,6 +160,9 @@ private[kafka010] class KafkaOffsetReaderConsumer(
       case SpecificTimestampRangeLimit(partitionTimestamps) =>
         fetchSpecificTimestampBasedOffsets(partitionTimestamps,
           failsOnNoMatchingOffset = isStartingOffsets).partitionToOffsets
+      case GlobalTimestampRangeLimit(timestamp) =>
+        fetchGlobalTimestampBasedOffsets(timestamp,
+          failsOnNoMatchingOffset = isStartingOffsets).partitionToOffsets
     }
   }
 
@@ -234,6 +237,39 @@ private[kafka010] class KafkaOffsetReaderConsumer(
       fnAssertFetchedOffsets)
   }
 
+  override def fetchGlobalTimestampBasedOffsets(
+      timestamp: Long,
+      failsOnNoMatchingOffset: Boolean): KafkaSourceOffset = {
+    val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = { partitions =>
+      logDebug(s"Partitions assigned to consumer: $partitions. Seeking to $timestamp")
+    }
+
+    val fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long] = { tps =>
+      val converted = tps.asScala.map(_ -> java.lang.Long.valueOf(timestamp)).toMap.asJava
+
+      val offsetForTime: ju.Map[TopicPartition, OffsetAndTimestamp] =
+        consumer.offsetsForTimes(converted)
+
+      offsetForTime.asScala.map { case (tp, offsetAndTimestamp) =>
+        if (failsOnNoMatchingOffset) {
+          assert(offsetAndTimestamp != null, "No offset matched from request of " +
+            s"topic-partition $tp and timestamp $timestamp.")
+        }
+
+        if (offsetAndTimestamp == null) {
+          tp -> KafkaOffsetRangeLimit.LATEST
+        } else {
+          tp -> offsetAndTimestamp.offset()
+        }
+      }.toMap
+    }
+
+    val fnAssertFetchedOffsets: Map[TopicPartition, Long] => Unit = { _ => }
+
+    fetchSpecificOffsets0(fnAssertParametersWithPartitions, fnRetrievePartitionOffsets,
+      fnAssertFetchedOffsets)
+  }
+
   private def fetchSpecificOffsets0(
       fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit,
       fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long],

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala

@@ -104,6 +104,8 @@ private[kafka010] class KafkaSource(
         case SpecificOffsetRangeLimit(p) => kafkaReader.fetchSpecificOffsets(p, reportDataLoss)
         case SpecificTimestampRangeLimit(p) =>
           kafkaReader.fetchSpecificTimestampBasedOffsets(p, failsOnNoMatchingOffset = true)
+        case GlobalTimestampRangeLimit(ts) =>
+          kafkaReader.fetchGlobalTimestampBasedOffsets(ts, failsOnNoMatchingOffset = true)
       }
       metadataLog.add(0, offsets)
       logInfo(s"Initial offsets: $offsets")