用火花窗函数计算移动平均时丢弃前几个值

Question

我试图计算按名称分组的列的季度移动平均线，我定义了一个火花窗口函数规范为

val wSpec1 = Window.partitionBy("name").orderBy("date").rowsBetween(-2, 0)

我的数据frame如下所示：

+-----+----------+-----------+------------------+
| name|      date|amountSpent|         movingAvg|
+-----+----------+-----------+------------------+
|  Bob|2016-01-01|       25.0|              25.0|
|  Bob|2016-02-02|       25.0|              25.0|
|  Bob|2016-03-03|       25.0|              25.0|
|  Bob|2016-04-04|       29.0|26.333333333333332|
|  Bob|2016-05-06|       27.0|              27.0|
|Alice|2016-01-01|       50.0|              50.0|
|Alice|2016-02-03|       45.0|              47.5|
|Alice|2016-03-04|       55.0|              50.0|
|Alice|2016-04-05|       60.0|53.333333333333336|
|Alice|2016-05-06|       65.0|              60.0|
+-----+----------+-----------+------------------+

import com.datastax.spark.connector._
import org.apache.spark.SparkContext
import org.apache.spark.SparkContext._
import org.apache.spark.SparkConf
import org.apache.spark.sql._
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.types.StructField
import org.apache.spark.sql.types.StringType
import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.types.DoubleType
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._
import org.apache.spark.sql.Row
import org.apache.spark.sql.types._
object Test {

  def main(args: Array[String]) {
    //val sparkSession = SparkSession.builder.master("local").appName("Test").config("spark.cassandra.connection.host", "localhost").config("spark.driver.host", "localhost").getOrCreate()
    val sparkSession = SparkSession.builder.master("local").appName("Test").config("spark.cassandra.connection.host", "localhost").config("spark.driver.host", "localhost").getOrCreate()
    val sc = sparkSession.sparkContext

    val sqlContext = new org.apache.spark.sql.SQLContext(sc)
    import sparkSession.implicits._

    val customers = sc.parallelize(List(("Alice", "2016-01-01", 50.00),
      ("Alice", "2016-02-03", 45.00),
      ("Alice", "2016-03-04", 55.00),
      ("Alice", "2016-04-05", 60.00),
      ("Alice", "2016-05-06", 65.00),
      ("Bob", "2016-01-01", 25.00),
      ("Bob", "2016-02-02", 25.00),
      ("Bob", "2016-03-03", 25.00),
      ("Bob", "2016-04-04", 29.00),
      ("Bob", "2016-05-06", 27.00))).toDF("name", "date", "amountSpent")

    import org.apache.spark.sql.expressions.Window
    import org.apache.spark.sql.functions._

    // Create a window spec.
    val wSpec1 = Window.partitionBy("name").orderBy("date").rowsBetween(-2, 0)

    val ls=customers.withColumn("movingAvg",avg(customers("amountSpent")).over(wSpec1))
    ls.show()

  }
}

Answer 1

我建议只计算平均值，如果窗口包含精确的3行（即跨越整个范围-2到0）

val ls=customers
.withColumn("count",count(($"amountSpent")).over(wSpec1))
.withColumn("movingAvg",when($"count"===3,avg(customers("amountSpent")).over(wSpec1)))

ls.show()


+-----+----------+-----------+-----+------------------+
| name|      date|amountSpent|count|         movingAvg|
+-----+----------+-----------+-----+------------------+
|  Bob|2016-01-01|       25.0|    1|              null|
|  Bob|2016-02-02|       25.0|    2|              null|
|  Bob|2016-03-03|       25.0|    3|              25.0|
|  Bob|2016-04-04|       29.0|    3|26.333333333333332|
|  Bob|2016-05-06|       27.0|    3|              27.0|
|Alice|2016-01-01|       50.0|    1|              null|
|Alice|2016-02-03|       45.0|    2|              null|
|Alice|2016-03-04|       55.0|    3|              50.0|
|Alice|2016-04-05|       60.0|    3|53.333333333333336|
|Alice|2016-05-06|       65.0|    3|              60.0|
+-----+----------+-----------+-----+------------------+