Add sequence/timestamp conversions and tests

Author: jpolchlo

src/main/scala/vectorpipe/sources/AugmentedDiffSource.scala

@@ -3,6 +3,8 @@ package vectorpipe.sources
 import java.io.{ByteArrayInputStream, File}
 import java.net.URI
 import java.nio.charset.StandardCharsets
+import java.sql.Timestamp
+import java.time.Instant
 import java.util.zip.GZIPInputStream
 
 import cats.implicits._
@@ -16,6 +18,8 @@ import io.circe.generic.auto._
 import io.circe.{yaml, _}
 import org.apache.commons.io.IOUtils
 import org.apache.spark.internal.Logging
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.functions.{floor, from_unixtime, to_timestamp, unix_timestamp}
 import org.joda.time.DateTime
 import vectorpipe.model.{AugmentedDiff, ElementWithSequence}
 
@@ -105,5 +109,17 @@ object AugmentedDiffSource extends Logging {
     }
   }
 
+  def timestampToSequence(timestamp: Timestamp): Int =
+    ((timestamp.toInstant.getEpochSecond - 1347432900) / 60).toInt
+
+  def timestampToSequence(timestamp: Column): Column =
+    floor((unix_timestamp(timestamp) - 1347432900) / 60)
+
+  def sequenceToTimestamp(sequence: Int): Timestamp =
+    Timestamp.from(Instant.ofEpochSecond(sequence.toLong * 60 + 1347432900L))
+
+  def sequenceToTimestamp(sequence: Column): Column =
+    to_timestamp(from_unixtime(sequence * 60 + 1347432900))
+
   case class State(last_run: DateTime, sequence: Int)
 }

src/test/scala/vectorpipe/sources/AugmentedDiffSourceTest.scala

@@ -0,0 +1,25 @@
+package vectorpipe.sources
+
+import org.apache.spark.sql.Row
+import org.scalatest.{FunSpec, Matchers}
+import vectorpipe.TestEnvironment
+
+class AugmentedDiffSourceSpec extends FunSpec with TestEnvironment with Matchers {
+
+  import ss.implicits._
+
+  describe("Timestamp to sequence conversion") {
+    it("should provide a round trip for simple conversion") {
+      AugmentedDiffSource.timestampToSequence(AugmentedDiffSource.sequenceToTimestamp(3700047)) should be (3700047)
+    }
+
+    it("should provide a round trip for column functions") {
+      val df = ss.createDataset(Seq(3700047)).toDF
+      (df.select(AugmentedDiffSource.sequenceToTimestamp('value) as 'time)
+         .select(AugmentedDiffSource.timestampToSequence('time) as 'value)
+         .first
+         .getLong(0)) should be (3700047)
+    }
+  }
+
+}