Dataframe 如何按列值在pyspark df中添加更多行

我被这个问题困扰了很长一段时间，可能使它比实际情况更严重。我会尽量简化它

我在代码中使用了pyspark和数据帧函数

我已经有一个df as：

+--+-----+---------+
|id|col1 |col2     |
+--+-----+---------+
|1 |Hello|Repeat   |
|2 |Word |Repeat   |
|3 |Aux  |No repeat|
|4 |Test |Repeat   |
+--+-----+---------+

我想要实现的是，当col2为“repeat”时，重复df的行，并在value+1中增加col1的值

+--+-----+---------+------+
|id|col1 |col2     |col3  |
+--+-----+---------+------+
|1 |Hello|Repeat   |Hello1|
|1 |Hello|Repeat   |Hello2|
|1 |Hello|Repeat   |Hello3|
|2 |Word |Repeat   |Word1 |
|2 |Word |Repeat   |Word2 |
|2 |Word |Repeat   |Word3 |
|3 |Aux  |No repeat|Aux   |
|4 |Test |Repeat   |Test1 |
|4 |Test |Repeat   |Test2 |
|4 |Test |Repeat   |Test3 |
+--+-----+---------+------+

我的第一种方法是使用withColumn操作符在udf的帮助下创建一个新列：

my_func = udf(lambda words: (words + str(i + 1 for i in range(3))), StringType())
df = df\
    .withColumn('col3', when(col('col2') == 'No Repeat', col('col1'))
                            .otherwise(my_func(col('col1'))))

但是，当我在df.show10中计算这个值时，它是错误的。我的猜测是因为我无法用withColumn函数以这种方式创建更多的行

因此，我决定采用另一种方法，但也没有成功。使用rdd.flatMap：

但是这里我丢失了df模式，我不能在else条件下抛出整行，它只抛出col1单词加上它的迭代器

你知道解决这个问题的正确方法吗

最后，我的问题是，我没有找到一种基于列值创建更多行的正确方法，因为我在这个世界上是新手。我发现的答案似乎也不适合这个问题

---

感谢所有帮助。

一种方法是使用条件并分配数组，然后分解

import pyspark.sql.functions as F

(df.withColumn("test",F.when(df['col2']=='Repeat',
       F.array([F.lit(str(i)) for i in range(1,4)])).otherwise(F.array(F.lit(''))))
  .withColumn("col3",F.explode(F.col("test"))).drop("test")
  .withColumn("col3",F.concat(F.col("col1"),F.col("col3")))).show()

@MohammadMurtazaHashmi建议的更简洁版本如下所示：

(df.withColumn("test",F.when(df['col2']=='Repeat',
     F.array([F.concat(F.col("col1"),F.lit(str(i))) for i in range(1,4)]))
    .otherwise(F.array(F.col("col1"))))
    .select("id","col1","col2", F.explode("test"))).show()

---

一种方法是使用条件并指定数组，然后分解

import pyspark.sql.functions as F

(df.withColumn("test",F.when(df['col2']=='Repeat',
       F.array([F.lit(str(i)) for i in range(1,4)])).otherwise(F.array(F.lit(''))))
  .withColumn("col3",F.explode(F.col("test"))).drop("test")
  .withColumn("col3",F.concat(F.col("col1"),F.col("col3")))).show()

@MohammadMurtazaHashmi建议的更简洁版本如下所示：

(df.withColumn("test",F.when(df['col2']=='Repeat',
     F.array([F.concat(F.col("col1"),F.lit(str(i))) for i in range(1,4)]))
    .otherwise(F.array(F.col("col1"))))
    .select("id","col1","col2", F.explode("test"))).show()

---

也许看一看DataFrame.explode方法？我在和rdd.flatmap进行斗争，因为它的性能很好，但是是的explode方法很好。也许看一看DataFrame.explode方法？我在和rdd.flatmap进行斗争，因为它的性能很好，但是是的explode方法很好。非常感谢您的回答@anky。最后，我采用了与你非常相似的方法。由于性能原因，我一直在避免使用explode函数，但刚才explode应该可以。@anky您也可以像这样使用F.array[F.concatF.colcol1，F.litstri for I in range 1,4]@MohammadMurtazaHashmi这确实更整洁。我想知道这样的方法应该存在，因为我熟悉使用熊猫的操作，但不知道在pyspark中如何操作。谢谢你启发我：非常感谢你的回答@anky。最后，我采用了与你非常相似的方法。由于性能原因，我一直在避免使用explode函数，但刚才explode应该可以。@anky您也可以像这样使用F.array[F.concatF.colcol1，F.litstri for I in range 1,4]@MohammadMurtazaHashmi这确实更整洁。我想知道这样的方法应该存在，因为我熟悉使用熊猫的操作，但不知道在pyspark中如何操作。谢谢你的启发：