如何根据pyspark中的条件设置新标志?

时间:2018-01-31 14:03:13

标签: python apache-spark pyspark spark-dataframe pyspark-sql

我有两个数据框,如下所示。

  df = spark.createDataFrame(sc.parallelize([[1,1,2],[1,2,9], [2,1,2],[2,2,1],
                [4,1,5],[4,2,6]]), ["sid","cid","Cr"])
  df.show() 
+---+---+---+
|sid|cid| Cr|
+---+---+---+
|  1|  1|  2|
|  1|  2|  9|
|  2|  1|  2|
|  2|  2|  1|
|  4|  1|  5|
|  4|  2|  6|
|  5|  1|  3|
|  5|  2|  8|
+---+---+---+

接下来我创建了df1,如下所示。

df1 = spark.createDataFrame(sc.parallelize([[1,1],[1,2],[1,3], [2,1],[2,2],[2,3],[4,1],[4,2],[4,3],[5,1],[5,2],[5,3]]), ["sid","cid"])

df1.show()
+---+---+
|sid|cid|
+---+---+
|  1|  1|
|  1|  2|
|  1|  3|
|  2|  1|
|  2|  2|
|  2|  3|
|  4|  1|
|  4|  2|
|  4|  3|
|  5|  1|
|  5|  2|
|  5|  3|
+---+---+

现在我希望我的最终输出应该如下所示。如果提供任何数据, 如果(df1.sid == df.sid)&(df1.cid == df.cid)则标记值1,否则为0。 缺少Cr值将为'0'

+---+---+---+----+
|sid|cid| Cr|flag|
+---+---+---+----+
|  1|  1|  2|  1 |
|  1|  2|  9|  1 |
|  1|  3|  0|  0 |
|  2|  1|  2|  1 |
|  2|  2|  1|  1 |
|  2|  3|  0|  0 |
|  4|  1|  5|  1 |
|  4|  2|  6|  1 |
|  4|  3|  0|  0 |
|  5|  1|  3|  1 |
|  5|  2|  8|  1 |
|  5|  3|  0|  0 |
+---+---+---+----+

请帮助我。

1 个答案:

答案 0 :(得分:1)

使用数据:

from pyspark.sql.functions import col, when, lit, coalesce

df = spark.createDataFrame(
    [(1, 1, 2), (1, 2, 9), (2, 1, 2), (2, 2, 1), (4, 1, 5), (4, 2, 6), (5, 1, 3), (5, 2, 8)], 
   ("sid", "cid", "Cr"))

df1 = spark.createDataFrame(
    [[1,1],[1,2],[1,3], [2,1],[2,2],[2,3],[4,1],[4,2],[4,3],[5,1],[5,2],[5,3]],
    ["sid","cid"])

外部联接:

joined = (df.alias("df")
    .join(
        df1.alias("df1"),
        (col("df.sid") == col("df1.sid")) & (col("df.cid") == col("df1.cid")),
        "rightouter"))

并选择

joined.select(
    col("df1.*"),
    coalesce(col("Cr"), lit(0)).alias("Cr"), 
    col("df.sid").isNotNull().cast("integer").alias("flag")
).orderBy("sid", "cid").show()

# +---+---+---+----+                 
# |sid|cid| Cr|flag|
# +---+---+---+----+
# |  1|  1|  2|   1|
# |  1|  2|  9|   1|
# |  1|  3|  0|   0|
# |  2|  1|  2|   1|
# |  2|  2|  1|   1|
# |  2|  3|  0|   0|
# |  4|  1|  5|   1|
# |  4|  2|  6|   1|
# |  4|  3|  0|   0|
# |  5|  1|  3|   1|
# |  5|  2|  8|   1|
# |  5|  3|  0|   0|
# +---+---+---+----+