如何根据原始数据帧中的总行数将数据帧拆分为两个数据帧

Posted 2023-04-17

【中文标题】如何根据原始数据帧中的总行数将数据帧拆分为两个数据帧

【英文标题】：How to split a dataframe in two dataframes based on the total number of rows in the original dataframe

【发布时间】：2020-08-13 16:11:13

【问题描述】：

您好，我是 spark 和 scala 的新手，我想拆分以下数据框：

df:
+----------+-----+------+----------+--------+
|        Ts| Temp|  Wind|  Precipit|Humidity|
+----------+-----+------+----------+--------+
|1579647600|   10|    22|        10|      50|
|1579734000|   11|    21|        10|      55|
|1579820400|   10|    18|        15|      60|
|1579906800|    9|    23|        20|      60|
|1579993200|    8|    24|        25|      50|
|1580079600|   10|    18|        27|      60|
|1580166000|   11|    20|        30|      50|
|1580252400|   12|    17|        15|      50|
|1580338800|   10|    14|        21|      50|
|1580425200|    9|    16|        25|      60|
-----------+-----+------+----------+--------+

生成的数据帧应如下所示：

df1:
+----------+-----+------+----------+--------+
|        Ts| Temp|  Wind|  Precipit|Humidity|
+----------+-----+------+----------+--------+
|1579647600|   10|    22|        10|      50|
|1579734000|   11|    21|        10|      55|
|1579820400|   10|    18|        15|      60|
|1579906800|    9|    23|        20|      60|
|1579993200|    8|    24|        25|      50|
|1580079600|   10|    18|        27|      60|
|1580166000|   11|    20|        30|      50|
|1580252400|   12|    17|        15|      50|
+----------+-----+------+----------+--------+
df2:
+----------+-----+------+----------+--------+
|        Ts| Temp|  Wind|  Precipit|Humidity|
+----------+-----+------+----------+--------+
|1580338800|   10|    14|        21|      50|
|1580425200|    9|    16|        25|      60|
-----------+-----+------+----------+--------+

其中 df1 占 df 顶行的 80%，而 df2 占 20%。

【参考方案1】：

尝试使用 monotonically_increasing_id() 函数和 window percent_rank()，因为此函数会保留顺序。

Example:

val df=sc.parallelize(Seq((1579647600,10,22,10,50),
(1579734000,11,21,10,55),
(1579820400,10,18,15,60),
(1579906800, 9,23,20,60),
(1579993200, 8,24,25,50),
(1580079600,10,18,27,60),
(1580166000,11,20,30,50),
(1580252400,12,17,15,50),
(1580338800,10,14,21,50),
(1580425200, 9,16,25,60)),10).toDF("Ts","Temp","Wind","Precipit","Humidity")

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

val df1=df.withColumn("mid",monotonically_increasing_id)
val df_above_80=df1.withColumn("pr",percent_rank().over(w)).filter(col("pr") >= 0.8).drop(Seq("mid","pr"):_*)
val df_below_80=df1.withColumn("pr",percent_rank().over(w)).filter(col("pr") < 0.8).drop(Seq("mid","pr"):_*)

df_below_80.show()
/*
+----------+----+----+--------+--------+
|        Ts|Temp|Wind|Precipit|Humidity|
+----------+----+----+--------+--------+
|1579647600|  10|  22|      10|      50|
|1579734000|  11|  21|      10|      55|
|1579820400|  10|  18|      15|      60|
|1579906800|   9|  23|      20|      60|
|1579993200|   8|  24|      25|      50|
|1580079600|  10|  18|      27|      60|
|1580166000|  11|  20|      30|      50|
|1580252400|  12|  17|      15|      50|
+----------+----+----+--------+--------+
*/

df_above_80.show()
/*
+----------+----+----+--------+--------+
|        Ts|Temp|Wind|Precipit|Humidity|
+----------+----+----+--------+--------+
|1580338800|  10|  14|      21|      50|
|1580425200|   9|  16|      25|      60|
+----------+----+----+--------+--------+
*/

【参考方案2】：

假设数据是随机拆分的：

val Array(df1, df2) = df.randomSplit(Array(0.8, 0.2))

但是，如果“顶行”是指示例数据框中的“Ts”列，那么您可以这样做：

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.col,percent_rank

val window = Window.partitionBy().orderBy(df['Ts'].desc())

val df1 = df.select('*', percent_rank().over(window).alias('rank')) 
  .filter(col('rank') >= 0.2) 
  .show()

val df2 = df.select('*', percent_rank().over(window).alias('rank')) 
  .filter(col('rank') < 0.2) 
  .show()

【讨论】：

我的意思是我想拆分数据帧，而不改变其行的顺序，首先出现的 80% 的行应该出现在 df1 中，剩下的 20% 的行应该出现在 df2 中跨度>

所以第二个选项将起作用，因为您的数据集按 Ts 排序。上面 484 的回答也对你有用