update

Author: Tyson Condie

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

@@ -58,11 +58,7 @@ private[kafka010] object KafkaSourceOffset {
 
   /**
    * Returns [[KafkaSourceOffset]] from a JSON [[SerializedOffset]]
-   * or [[KafkaSourceOffset]]
    */
-  def apply(offset: Offset): KafkaSourceOffset = offset match {
-    case kso: KafkaSourceOffset => KafkaSourceOffset(kso.partitionToOffsets)
-    case so: SerializedOffset =>
-      KafkaSourceOffset(JsonUtils.partitionOffsets(offset.json))
-  }
+  def apply(offset: SerializedOffset): KafkaSourceOffset =
+    KafkaSourceOffset(JsonUtils.partitionOffsets(offset.json))
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/FileStreamSource.scala

@@ -131,8 +131,8 @@ class FileStreamSource(
    * Returns the data that is between the offsets (`start`, `end`].
    */
   override def getBatch(start: Option[Offset], end: Offset): DataFrame = {
-    val startId = start.map(LongOffset(_).offset).getOrElse(-1L)
-    val endId = LongOffset(end).offset
+    val startId = start.map(LongOffset.convert(_).offset).getOrElse(-1L)
+    val endId = LongOffset.convert(end).offset
 
     assert(startId <= endId)
     val files = metadataLog.get(Some(startId + 1), Some(endId)).flatMap(_._2)

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/LongOffset.scala

@@ -33,13 +33,17 @@ case class LongOffset(offset: Long) extends Offset {
 object LongOffset {
 
   /**
-   * LongOffset factory from Offset.
-   * @param offset used to initialize LongOffset. Assumed to be
-   *               either a LongOffset or SerializedOffset
+   * LongOffset factory from serialized offset.
    * @return new LongOffset
    */
-  def apply(offset: Offset) : LongOffset = offset match {
-    case lo: LongOffset => new LongOffset(lo.offset)
-    case so: SerializedOffset => new LongOffset(offset.json.toLong)
+  def apply(offset: SerializedOffset) : LongOffset = new LongOffset(offset.json.toLong)
+
+  /**
+   * Convert generic Offset to LongOffset if possible.
+   * @return converted LongOffset
+   */
+  def convert(offset: Offset): LongOffset = offset match {
+    case lo: LongOffset => lo
+    case so: SerializedOffset => LongOffset(so)
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/memory.scala

@@ -106,8 +106,8 @@ case class MemoryStream[A : Encoder](id: Int, sqlContext: SQLContext)
   override def getBatch(start: Option[Offset], end: Offset): DataFrame = {
     // Compute the internal batch numbers to fetch: [startOrdinal, endOrdinal)
     val startOrdinal =
-      start.map(LongOffset(_)).getOrElse(LongOffset(-1)).offset.toInt + 1
-    val endOrdinal = LongOffset(end).offset.toInt + 1
+      start.map(LongOffset.convert(_)).getOrElse(LongOffset(-1)).offset.toInt + 1
+    val endOrdinal = LongOffset.convert(end).offset.toInt + 1
 
     // Internal buffer only holds the batches after lastCommittedOffset.
     val newBlocks = synchronized {