根据列的值添加火花数据集中的行号

Question

我有一个示例Dataset<Row> 如下

+----+----+
|col1|col2|
+----+----+
|   a|  a |
|   a|  b |
|   b|  c |
|   c|  d |
+----+----+

基于映射值 Map sample ={"key1" = 2,"key2"=4} 我必须为每一行添加一个 row_number 以便 where col1 or col2 = 'a' 我必须添加 2 个连续的行号和其他作为 4基于来自map 的值，结果应如下所示。 无论col1 or col2 ='a' 行在哪里复制了两次 次，如果不是，它复制了四次 次，后面是数字

+----+----+----------+
|col1|col2|limit_fact|
+----+----+----------+
|   a|  a |         1|
|   a|  a |         2|
|   a|  b |         3|
|   a|  b |         4|
|   b|  c |         5|
|   b|  c |         6|
|   b|  c |         7|
|   b|  c |         8|
|   c|  d |         9|
|   c|  d |        10|
|   c|  d |        11|
|   c|  d |        12|
+----+----+----------+

Answer 1

scala> import org.apache.spark.sql.expressions.{UserDefinedFunction, Window}

scala> val key1 = 2

scala> val key2 = 4

scala> def repeatudf:UserDefinedFunction = udf((key:Int) => {var out = "1"
     | val rng = 1 to (key-1)
     | for (i <- rng) {out = out + "," + "1"}
     | out
     | })

scala> df.show
+----+----+
|col1|col2|
+----+----+
|   a|   a|
|   a|   b|
|   b|   c|
|   c|   d|
+----+----+


scala> val df1 =  df.withColumn("map", when(col("col1") === "a" || col("col2") === "a",  repeatudf(lit(key1))).otherwise(repeatudf(lit(key2))))
                    .withColumn("map", explode(split(col("map"), ",")))
                    .drop("map")

scala> df1.show()
+----+----+
|col1|col2|
+----+----+
|   a|   a|
|   a|   a|
|   a|   b|
|   a|   b|
|   b|   c|
|   b|   c|
|   b|   c|
|   b|   c|
|   c|   d|
|   c|   d|
|   c|   d|
|   c|   d|
+----+----+

scala> val w =  Window.orderBy(col("col1"), col("col2"))

scala> df1.withColumn("limit_fact", row_number.over(w)).show
+----+----+----------+
|col1|col2|limit_fact|
+----+----+----------+
|   a|   a|         1|
|   a|   a|         2|
|   a|   b|         3|
|   a|   b|         4|
|   b|   c|         5|
|   b|   c|         6|
|   b|   c|         7|
|   b|   c|         8|
|   c|   d|         9|
|   c|   d|        10|
|   c|   d|        11|
|   c|   d|        12|
+----+----+----------+