title: Grouping by an alternating sequence in Spark
body: <p>I have a set of data where I can identify an alternating sequence. However, I want to group this data into one chunk leaving all the other data unchanged. That is, where ever flickering is occurring in id, I want to overwrite that group of id with the single id that makes since to the order. As a small example, consider</p>
  
  <pre><code>val dataDF = Seq(
      ("a", "silom", 3, 1),
      ("a", "silom", 2, 2),
      ("a", "silom", 1, 3),
      ("a", "silom", 0, 4),  // flickering; id=0
      ("a", "silom", 1, 5),  // flickering; id=0
      ("a", "silom", 0, 6),  // flickering; id=0
      ("a", "silom", 1, 7),
      ("a", "silom", 2, 8),
      ("a", "silom", 3, 9),
      ("a", "silom", 4, 10),
      ("a", "silom", 3, 11),  // flickering and so on
      ("a", "silom", 4, 12),
      ("a", "silom", 3, 13),
      ("a", "silom", 4, 14),
      ("a", "silom", 5, 15)
  ).toDF("user", "cat", "id", "time_sec")
  
  val resultDataDF = Seq(
      ("a", "silom", 3, 1),
      ("a", "silom", 2, 2),
      ("a", "silom", 1, 3),
      ("a", "silom", 0, 15),  // grouped by flickering summing on time_sec
      ("a", "silom", 1, 7),
      ("a", "silom", 2, 8),
      ("a", "silom", 3, 9),
      ("a", "silom", 4, 60),
      ("a", "silom", 5, 15). // grouped by flickering summing on time_sec
  ).toDF("user", "cat", "id", "time_sec")
  </code></pre>
  
  <p>Now a more realistic MWE. In this case, we can have multiple users and cat; unfortunately, this approach doesnt use the dataframe API and needs to collect data to the driver. This isn't scalable and needs to recursively call <code>getGrps</code> by dropping the length of the returned array indices.</p>
  
  <blockquote>
    <p>How can I implement this using the dataframe API so as not to need collect the data to the driver which would be impossible due to size? Also, if there is a better way to do this, what would that be?</p>
  </blockquote>
  
  <pre><code>import org.apache.spark.sql.functions._
  import org.apache.spark.sql.expressions.Window
  import org.apache.spark.sql.types.DoubleType
  import scala.collection.mutable.WrappedArray
  
  val dataDF = Seq(
      ("a", "silom", 3, 1),
      ("a", "silom", 2, 2),
      ("a", "silom", 1, 3),
      ("a", "silom", 0, 4),
      ("a", "silom", 1, 5),
      ("a", "silom", 0, 6),
      ("a", "silom", 1, 7),
      ("a", "silom", 2, 8),
      ("a", "silom", 3, 9),
      ("a", "silom", 4, 10),
      ("a", "silom", 3, 11),
      ("a", "silom", 4, 12),
      ("a", "silom", 3, 13),
      ("a", "silom", 4, 14),
      ("a", "silom", 5, 15),
      ("a", "suk", 18, 1),
      ("a", "suk", 19, 2),
      ("a", "suk", 20, 3),
      ("a", "suk", 21, 4),
      ("a", "suk", 20, 5),
      ("a", "suk", 21, 6),
      ("a", "suk", 0, 7),
      ("a", "suk", 1, 8),
      ("a", "suk", 2, 9),
      ("a", "suk", 3, 10),
      ("a", "suk", 4, 11),
      ("a", "suk", 3, 12),
      ("a", "suk", 4, 13),
      ("a", "suk", 3, 14),
      ("a", "suk", 5, 15),
      ("b", "silom", 4, 1),
      ("b", "silom", 3, 2),
      ("b", "silom", 2, 3),
      ("b", "silom", 1, 4),
      ("b", "silom", 0, 5),
      ("b", "silom", 1, 6),
      ("b", "silom", 2, 7),
      ("b", "silom", 3, 8),
      ("b", "silom", 4, 9),
      ("b", "silom", 5, 10),
      ("b", "silom", 6, 11),
      ("b", "silom", 7, 12),
      ("b", "silom", 8, 13),
      ("b", "silom", 9, 14),
      ("b", "silom", 10, 15),
      ("b", "suk", 11, 1),
      ("b", "suk", 12, 2),
      ("b", "suk", 13, 3),
      ("b", "suk", 14, 4),
      ("b", "suk", 13, 5),
      ("b", "suk", 14, 6),
      ("b", "suk", 13, 7),
      ("b", "suk", 12, 8),
      ("b", "suk", 11, 9),
      ("b", "suk", 10, 10),
      ("b", "suk", 9, 11),
      ("b", "suk", 8, 12),
      ("b", "suk", 7, 13),
      ("b", "suk", 6, 14),
      ("b", "suk", 5, 15)
  ).toDF("user", "cat", "id", "time_sec")
  val recastDataDF = dataDF.withColumn("id", $"id".cast(DoubleType))
  
  val category = recastDataDF.select("cat").distinct.collect.map(x =&gt; x(0).toString)
  
  val data = recastDataDF
      .select($"*" +: category.map(
          name =&gt; 
          lag("id", 1).over(
              Window.partitionBy("user", "cat").orderBy("time_sec")
          )
          .alias(s"lag_${name}_id")): _*)
      .withColumn("sequencing_diff", when($"cat" === "silom", ($"lag_silom_id" - $"id").cast(DoubleType))
                  .otherwise(($"lag_suk_id" - $"id")))
      .drop("lag_silom_id", "lag_suk_id")
      .withColumn("rn", row_number.over(Window.partitionBy("user", "cat").orderBy("time_sec")).cast(DoubleType))
      .withColumn("zipped", array("user", "cat", "sequencing_diff", "rn", "id"))
  
  // non dataframe API approach (not scalable)
  // needs to collect data to driver to process
  val iterTuples = data.select("zipped").collect.map(x =&gt; x(0).asInstanceOf[WrappedArray[Any]]).map(x =&gt; x.toArray)
  
  val shifted: Array[Array[Any]] = iterTuples.drop(1)
  val combined = iterTuples
      .zipAll(shifted, Array("", "", Double.NaN, Double.NaN, Double.NaN), Array("", "", Double.NaN, Double.NaN, Double.NaN))
  
  val testArr = combined.map{
      case (data0, data1) =&gt;
      if(data1(3).toString.toDouble &gt; 2 &amp;&amp; data0(3).toString.toDouble &gt; 2 &amp;&amp; data1(0) == data0(0) &amp;&amp; data1(1) == data0(1)) {
          if(data0(2) != data1(2) &amp;&amp; data0(2).toString.toDouble + data1(2).toString.toDouble == 0) {
              (data1(0), data1(1), data1(3), data0(4))
          }
          else ("", "", Double.NaN, Double.NaN)
      }
      else ("", "", Double.NaN, Double.NaN)
  }
      .filter(t =&gt; t._1 != "" &amp;&amp; t._2 != "" &amp;&amp; t._3 == t._3 &amp;&amp; t._4 == t._4)  // fast NaN removal
  
  val typeMappedArray = testArr.map(x =&gt; (x._1.toString, x._2.toString, x._3.toString.toDouble, x._4.toString.toDouble))
  
  def getGrps(arr: Array[(String, String, Double, Double)]): (Array[Double], Double, String, String) = {
  
      if(arr.nonEmpty) {
          val user = arr.take(1)(0)._1
          val cat = arr.take(1)(0)._2
          val rowNum = arr.take(1)(0)._3
          val keepID = arr.take(1)(0)._4
          val newArr = arr.drop(1)
  
          val rowNums = (Array(rowNum)) ++ newArr.zipWithIndex.map{
              case (tups, idx) =&gt; 
              if(rowNum + idx + 1 == tups._3) {
                  rowNum + 1 + idx
              }
              else Double.NaN
          }
              .filter(v =&gt; v == v)
  
          (rowNums, keepID, user, cat)
      }
      else (Array(Double.NaN), Double.NaN, "", "")
  }
  
  
  // after overwriting, this would allow me to group by user, cat, id to sum the time
  getGrps(typeMappedArray)  // returns rows number to overwrite, value to overwrite id with, user, cat
  res0: (Array(5.0, 6.0, 7.0),0.0,a,silom)
  
  getGrps(typeMappedArray.drop(3))
  res1: (Array(11.0, 12.0, 13.0, 14.0),4.0,a,silom)
  </code></pre>
  
  <hr>
  
  <p>A second approach using <code>collect_list</code> but this relies on <code>getGrps</code> working recursively which I cannot get working properly. Here is the code I have so far with a modified <code>getGrps</code> for the the <code>collect_list</code> minus recursive.</p>
  
  <pre><code>val data = recastDataDF
      .select($"*" +: category.map(
          name =&gt; 
          lag("id", 1).over(
              Window.partitionBy("user", "cat").orderBy("time_sec")
          )
          .alias(s"lag_${name}_id")): _*)
      .withColumn("sequencing_diff", when($"cat" === "silom", ($"lag_silom_id" - $"id").cast(DoubleType))
                  .otherwise(($"lag_suk_id" - $"id")))
      .drop("lag_silom_id", "lag_suk_id")
      .withColumn("rn", row_number.over(Window.partitionBy("user", "cat").orderBy("time_sec")).cast(DoubleType))
      .withColumn("id_rn", array($"id", $"rn", $"sequencing_diff"))
      .groupBy($"user", $"cat").agg(collect_list($"id_rn").alias("array_data"))
  
  // collect one row to develop how the UDF would work
  val testList = data.where($"user" === "a" &amp;&amp; $"cat" === "silom").select("array_data").collect
      .map(x =&gt; x(0).asInstanceOf[WrappedArray[WrappedArray[Any]]])
      .map(x =&gt; x.toArray)
      .head
      .map(x =&gt; (x(0).toString.toDouble, x(1).toString.toDouble, x(2).asInstanceOf[Double]))
  
  // this code would be in the UDF; that is, we would pass array_data to the UDF
  scala.util.Sorting.stableSort(testList, (e1: (Double, Double, Double), e2: (Double, Double, Double)) =&gt; e1._2 &lt; e2._2)
  
  val shifted: Array[(Double, Double, Double)] = testList.drop(1)
  val combined = testList
      .zipAll(shifted, (Double.NaN, Double.NaN, Double.NaN), (Double.NaN, Double.NaN, Double.NaN))
  
  val testArr = combined.map{
      case (data0, data1) =&gt;
      if(data0._3 != data1._3 &amp;&amp; data0._2 &gt; 1) {
          (data0._2, data0._1)
      }
      else (Double.NaN, Double.NaN)
      }
      .filter(t =&gt; t._1 == t._1 &amp;&amp; t._1 == t._1) 
  
  // called inside the UDF
  def getGrps2(arr: Array[(Double, Double)]): (Array[Double], Double) = {
      // no need for user or cat
  
      if(arr.nonEmpty) {
          val rowNum = arr.take(1)(0)._1
          val keepID = arr.take(1)(0)._2
          val newArr = arr.drop(1)
  
          val rowNums = (Array(rowNum)) ++ newArr.zipWithIndex.map{
              case (tups, idx) =&gt; 
              if(rowNum + idx + 1 == tups._1) {
                  rowNum + 1 + idx
              }
              else Double.NaN
          }
              .filter(v =&gt; v == v)
  
          (rowNums, keepID)
      }
      else (Array(Double.NaN), Double.NaN)
  }
  </code></pre>
  
  <p>We would <code>.withColumn("data_to_update", udf)</code> and the <code>data_to_update</code> column would be a <code>WrappedArray[Tuple2[Array[Double], Double]]</code> with row_numbers to id to overwrite. The result for user <code>a</code>, cat <code>silom</code> would be</p>
  
  <p><code>WrappedArray((Array(4.0, 5.0, 6.0),0.0), (Array(10.0, 11.0, 12.0, 13.0),4.0))</code></p>
  
  <p>The array pieces are row numbers and the Double is id to update those rows with</p>