Scala - 將函數應用於數據框列中的每個值

Question

我有一個函數，它采用LocalDate （它可以采用任何其他類型）並返回一個DataFrame ，例如：

def genDataFrame(refDate: LocalDate): DataFrame = {
  Seq(
    (refDate,refDate.minusDays(7)),
    (refDate.plusDays(3),refDate.plusDays(7))
  ).toDF("col_A","col_B")
}

genDataFrame(LocalDate.parse("2021-07-02"))輸出：

+----------+----------+
|     col_A|     col_B|
+----------+----------+
|2021-07-02|2021-06-25|
|2021-07-05|2021-07-09|
+----------+----------+

我想將此函數應用於數據框列（顯然包含LocalDate值）中的每個元素，例如：

val myDate = LocalDate.parse("2021-07-02")

val df = Seq(
  (myDate),
  (myDate.plusDays(1)),
  (myDate.plusDays(3))
).toDF("date")

df ：

+----------+
|      date|
+----------+
|2021-07-02|
|2021-07-03|
|2021-07-05|
+----------+

所需輸出：

+----------+----------+
|     col_A|     col_B|
+----------+----------+
|2021-07-02|2021-06-25|
|2021-07-05|2021-07-09|
|2021-07-03|2021-06-26|
|2021-07-06|2021-07-10|
|2021-07-05|2021-06-28|
|2021-07-08|2021-07-12|
+----------+----------+

我怎么能做到這一點（不使用collect ）？

Answer 1

您始終可以將數據框轉換為延遲評估的視圖並使用 Spark SQL：

val df_2 = df.map(x => x.getDate(0).toLocalDate()).withColumnRenamed("value", "col_A")
.withColumn("col_B", col("col_A"))
df_2.createOrReplaceTempView("test")

有了它，您可以創建一個像這樣的視圖：

+----------+----------+
|     col_A|     col_B|
+----------+----------+
|2021-07-02|2021-07-02|
|2021-07-03|2021-07-03|
|2021-07-05|2021-07-05|
+----------+----------+

然后你可以使用我覺得更直觀的 SQL：

spark.sql(s"""SELECT col_A, date_add(col_B, -7) as col_B FROM test
UNION
SELECT date_add(col_A, 3), date_add(col_B, 7) as col_B FROM test""")
.show()

這將您的預期輸出作為 DataFrame 提供：

+----------+----------+
|     col_A|     col_B|
+----------+----------+
|2021-07-02|2021-06-25|
|2021-07-03|2021-06-26|
|2021-07-05|2021-06-28|
|2021-07-05|2021-07-09|
|2021-07-06|2021-07-10|
|2021-07-08|2021-07-12|
+----------+----------+