Создать новый столбец на основе значений строки scala spark

Question

У меня есть фрейм данных с несколькими столбцами. Я хочу создать новый столбец, основанный на следующей идее: если в строке есть один 0, поместите 0 в новый столбец, в противном случае введите 1

+-------+--------+-----+
|A      |B       |C    |
+-------+--------+-----+
|     1 |      1 |   1 | 
|     1 |      0 |   1 | 
+-------+--------+-----+

Желаемый результат:

+-------+--------+-----+-------+
|A      |B       |C    |New_Col|
+-------+--------+-----+-------+
|     1 |      1 |   1 |   1   |
|     1 |      0 |   1 |   0   |
+-------+--------+-----+-------+

Я пытался уменьшить и суммировать значения, а затем сравнить размеры и суммы, но не смог

Answer 1

Вы можете использовать функцию when для условного добавления столбца в кадре данных искры

val spark = SparkSession.builder().master("local[*]").appName("test").getOrCreate()

import spark.implicits._

val someDF = Seq(
      (1,1,1),
      (1,0,1),
      (1,1,1)).toDF("a","b","c")

someDF.withColumn("new",when($"a" + $"b" + $"c" ===lit(3),1).otherwise(0)).show()

Вот как будет выглядеть вывод

+---+---+---+---+
|  a|  b|  c|new|
+---+---+---+---+
|  1|  1|  1|  1|
|  1|  0|  1|  0|
|  1|  1|  1|  1|
+---+---+---+---+

Answer 2

Вы можете использовать when и otherwise, используя условие ИЛИ.

    val df = spark.createDataFrame(Seq((1,1,0), (1,2,3),(2,5,0))).toDF("a","b","c")
    df.withColumn("d", when(col("a") === 0 or col("b") === 0 or col("c") === 0, lit(0)).otherwise(lit(1))).show()

Вы также можете оценивать столбцы динамически на основе списка столбцов.

    val df = spark.createDataFrame(Seq((1,1,0), (1,2,3),(2,5,0))).toDF("a","b","c")

    val cols = Seq("a","b","c")
    val initCol = lit(0)===1
    val col1 = cols.foldLeft(initCol)((x,y) => x or col(y) ===0)

    df.withColumn("d", when(col1, lit(0)).otherwise(lit(1))).show()

Результат

+---+---+---+---+
|  a|  b|  c|  d|
+---+---+---+---+
|  1|  1|  0|  0|
|  1|  2|  3|  1|
|  2|  5|  0|  0|
+---+---+---+---+

Answer 3

Использование функции «В противном случае» может помочь. df.withColumn("NewCol", when(col("A") === 0,0).when(col("B") === 0,0).when(col("C") === 0,0).otherwise(1))