scala中的模拟火花列函数答案

【问题标题】：mock spark column functions in scalascala中的模拟火花列函数
【发布时间】：2021-05-08 05:41:12
【问题描述】：

我的代码使用monotonically_increasing_id函数是scala

val df = List(("oleg"), ("maxim")).toDF("first_name")
   .withColumn("row_id", monotonically_increasing_id)

我想在我的单元测试中模拟它，以便它返回整数 0, 1, 2, 3, ...

在我的 spark-shell 中，它返回所需的结果。

scala> df.show
+----------+------+
|first_name|row_id|
+----------+------+
|      oleg|     0|
|     maxim|     1|
+----------+------+

但在我的 scala 应用程序中，结果是不同的。

如何模拟列函数？

【问题讨论】：

我不确定我是否理解您的问题。您能否提供更多详细信息，例如您尝试测试的代码 sn-p？
@Oli 我编辑了我的答案
他们说Don't mock type you don't own - 这听起来像是一个很好的例子，为什么你不应该这样做 - 不能保证 monotonically_increasing_id 会返回连续的数字。
@user10958683 我看到了这个答案。它与嘲笑的感觉相矛盾：stackoverflow.com/questions/2665812/what-is-mocking。此外，我不在乎我的生产代码中的数字是否是连续的，只要它们是唯一的。我需要在单元测试中预测结果

标签： scala unit-testing apache-spark

【解决方案1】：

模拟这样一个函数以使其产生一个序列并不简单。实际上，spark 是一个并行计算引擎，因此按顺序访问数据很复杂。

这是一个你可以尝试的解决方案。

让我们定义一个压缩数据框的函数：

    def zip(df : DataFrame, name : String) = {
        df.withColumn(name, monotonically_increasing_id)
    }

那我们就默认使用这个zip函数重写我们要测试的函数吧：

    def fun(df : DataFrame,
            zipFun : (DataFrame, String) => DataFrame = zip) : DataFrame = {
        zipFun(df, "id_row")
    }
    // let 's see what it does
    fun(spark.range(5).toDF).show()
    +---+----------+
    | id|    id_row|
    +---+----------+
    |  0|         0|
    |  1|         1|
    |  2|8589934592|
    |  3|8589934593|
    |  4|8589934594|
    +---+----------+

和之前一样，让我们编写一个新函数，使用来自 RDD API 的zipWithIndex。这有点乏味，因为我们必须在两个 API 之间来回切换。

    def zip2(df : DataFrame, name : String) = {
        val rdd = df.rdd.zipWithIndex
            .map{ case (row, i) => Row.fromSeq(row.toSeq :+ i) }
        val newSchema = df.schema.add(StructField(name, LongType, false))
        df.sparkSession.createDataFrame(rdd, newSchema)
    }
    fun(spark.range(5).toDF, zip2)
    +---+------+
    | id|id_row|
    +---+------+
    |  0|     0|
    |  1|     1|
    |  2|     2|
    |  3|     3|
    |  4|     4|
    +---+------+

你可以调整zip2，例如将i乘以2，得到你想要的。

【讨论】：

谢谢，您可以查看我在单独答案中发布的解决方法

【解决方案2】：

根据@Oli 的回答，我想出了以下解决方法：

val df = List(("oleg"), ("maxim")).toDF("first_name")
   .withColumn("row_id", monotonically_increasing_id)
   .withColumn("test_id", row_number().over(Window.orderBy("row_id")))

它解决了我的问题，但我仍然对模拟列函数感兴趣。

【讨论】：

【解决方案3】：

我用这段代码模拟了我的 spark 函数：

val s = typedLit[Timestamp](Timestamp.valueOf("2021-05-07 15:00:46.394"))
implicit val ds = DefaultAnswer(CALLS_REAL_METHODS)
withObjectMocked[functions.type] {
when(functions.current_timestamp()).thenReturn(s)
        // spark logic
}

【讨论】：