python - 用现有数据框中的一些选定的行集形成一个新的spark数据框

Question

我有一个 Spark 数据帧df，行数为10 ^ 8。我在该数据框架上添加了一列作为rowId，我想用作主键。
我已经使用命令完成了与

df.withColumn(“rowId”，monotonically_increasing_id())

现在我要从该数据帧中选择一个新的数据帧，该数据帧具有一些选定的行数，这些行的索引以列表的形式对我来说是已知的。
如果有人可以帮助我使用列表中选定的行数来形成新的数据框，这对我会有所帮助。

Answer 1

我不确定我是否理解您的用例。 monotonically_increasing_id()可保证唯一的ID，但不能保证ID以0或1开头或数字是连续的。在下面的示例中，我对数据框进行了重新分区，以表明monotonically_increasing_id()不保证连续的id。无论如何，假设驱动程序内存中有所需的索引列表，则可以在添加索引列之后简单地加入数据框。

//Create a sample dataframe and add rowId column. 
//Note that you may see 0,1,2 as rowIds if you dont repartition. 

val df = List("A","B","C").toDF.repartition(5).withColumn("rowId", monotonically_increasing_id())
df.show()
+-----+------------+
|value|       rowId|
+-----+------------+
|    A|  8589934592|
|    B|  8589934593|
|    C| 34359738368|
+-----+------------+

//[Option 1] to join with indexes we need to add index column to our DataFrame. Assuming your indexes align with sorted rowId
val w = org.apache.spark.sql.expressions.Window.orderBy("rowId")
val result = df.withColumn("index", row_number().over(w) - 1)

//here is our indexes. let convert it to Dataframe to prepare for join
val indexes = List(0, 2).toDF
//finally join
result.join(indexes, result("index") === indexes("value")).show()
+-----+-----------+-----+-----+
|value|      rowId|index|value|
+-----+-----------+-----+-----+
|    A| 8589934592|    0|    0|
|    C|34359738368|    2|    2|
+-----+-----------+-----+-----+


//[Option 2] if your list is small and can easily be sent to all workers, you can also simply filter
result.filter(result("index").isin(List(0, 2):_*)).show()
+-----+-----------+-----+
|value|      rowId|index|
+-----+-----------+-----+
|    A| 8589934592|    0|
|    C|34359738368|    2|
+-----+-----------+-----+