如何将给定条件-python的Dataframe行拆分为两行

Question

我有一个Spark Dataframe如下，

+-------+-------+-----+
|  s1   |   s2  |isVal|
+-------+-------+-----+
|a      |aa     |    1|
|b      |bb     |    0|
|c      |cc     |    1|
|d      |dd     |    0|
|e      |ee     |    1|
+-------+-------+-----+

我想检查每行中的 isVal 值，如果该值等于 1 ，则该行应拆分为两行。例如：考虑上面数据帧的前两行，结果应该如下，

+-------+-------+
|  s1   |  isVal|
+-------+-------+
|a      |      1|
|aa     |      1|
|b      |      0|
+-------+-------+

请帮助使用python构建逻辑。我尝试使用flatmap构建逻辑，但它没有提供预期的结果。

Answer 1

这是解决方案。

>>> from pyspark.sql.functions import array, col, explode
>>>
>>> d = [{'s1': 'a', 's2': 'aa', 'isVal':1}, {'s1': 'b', 's2': 'bb', 'isVal':0}, {'s1': 'c', 's2': '
cc', 'isVal':1}, {'s1': 'd', 's2': 'dd', 'isVal':0}]
>>> df = spark.createDataFrame(d)
>>> df.show()
+-----+---+---+
|isVal| s1| s2|
+-----+---+---+
|    1|  a| aa|
|    0|  b| bb|
|    1|  c| cc|
|    0|  d| dd|
+-----+---+---+

>>> df1 = df.where(col("isVal")==1).withColumn("s3",array(col("s1"), col("s2"))).select(col("s3"), col("isVal")).withColumn("s1",explode(col("s3"))).drop(col("s3"))
>>> df0 = df.select(col("isVal"), col("s1")).where(col("isVal")==0)
>>> df2 = df1.union(df0)
>>> df2.show()
+-----+---+
|isVal| s1|
+-----+---+
|    1|  a|
|    1| aa|
|    1|  c|
|    1| cc|
|    0|  b|
|    0|  d|
+-----+---+