加入两个连接键不同的DataFrames，只选择一些列

Question

我想做的是：

连接两个 DataFrameA并B使用它们各自的id列a_id和b_id. 我想从中选择所有列A和两个特定列B

我尝试了类似下面用不同引号括起来的内容，但仍然无法正常工作。我觉得在 pyspark 中，应该有一个简单的方法来做到这一点。

A_B = A.join(B, A.id == B.id).select(A.*, B.b1, B.b2)

我知道你可以写

A_B = sqlContext.sql("SELECT A.*, B.b1, B.b2 FROM A JOIN B ON A.a_id = B.b_id")

这样做，但我想更像上面的伪代码。

Answer 1

你的伪代码基本上是正确的。如果该id列存在于两个 DataFrame 中，则此稍微修改的版本将起作用：

A_B = A.join(B, on="id").select("A.*", "B.b1", "B.b2")

从文档中pyspark.sql.DataFrame.join()：

如果on是指示连接列名称的字符串或字符串列表，则该列必须存在于两侧，这将执行等连接。

由于键不同，您可以使用withColumn()(或withColumnRenamed()) 在两个 DataFrame 中创建一个具有相同名称的列：

A_B = A.withColumn("id", col("a_id")).join(B.withColumn("id", col("b_id")), on="id")\
    .select("A.*", "B.b1", "B.b2")

如果你的 DataFrames 有很长很复杂的名字，你也可以用它alias()来简化事情：

A_B = long_data_frame_name1.alias("A").withColumn("id", col("a_id"))\
    .join(long_data_frame_name2.alias("B").withColumn("id", col("b_id")), on="id")\
    .select("A.*", "B.b1", "B.b2")