Scala Spark: лучший способ создать столбец для диапазонов

Question

Я ищу способ улучшить и уменьшить количество повторяющегося кода. Мне нужен один столбец для каждого диапазона. Есть ли способ создать метод, который может обрабатывать генерацию столбцов здесь?

df
  .withColumn(
    "daysbetween",
    fn.datediff($"date1", $"date2")
  )
  .withColumn(
    "0",
    fn.when(
      $"daysbetween" >= -30,
      $"TotalPrice"
    )
  )
  .withColumn(
    "-30",
    fn.when(
      $"daysbetween".between(-60, -31),
      $"TotalPrice"
    )
  )
  .withColumn(
    "-60",
    fn.when(
      $"daysbetween".between(-90, -61),
      $"TotalPrice"
    )
  )
  .withColumn(
    "<-90",
    fn.when(
      $"daysbetween" < -90,
      $"TotalPrice"
    )
  )

Это плохо в двух отношениях. Во-первых, у меня много повторяющегося кода. Во-вторых, я использую числа маги c, и эту проблему мне тоже нужно решить.

Буду признателен за помощь, ребята.

Привет Джо

Answer 1

Надеюсь, этот код поможет мне создать список столбцов, связывающих его с именем. НЕТ, если вы хотите автоматически сгенерировать логи генерации столбцов c

val dynamicCols : List[(String, Column)] = List(
      ("daysbetween",
         fn.datediff($"date1", $"date2")
      ),
      ("0", fn.when(
      $"daysbetween" >= -30,
      $"TotalPrice"
    )),
      ("-30", fn.when(
        $"daysbetween" >= -30,
        $"TotalPrice"
      )),
      ("-30",
        fn.when(
          $"daysbetween".between(-60, -31),
          $"TotalPrice"
        )),
      ("-60",
        $"".when(
          $"daysbetween".between(-90, -61),
          $"TotalPrice"
        )),
      ("<-90",
        fn.when(
          $"daysbetween" < -90,
          $"TotalPrice"
        ))
    )

    val df  = dynamicCols.foldLeft(df)((df : DataFrame,colInfo : (String,Column)) => {
      df.withColumn(colInfo._1,colInfo._2)
    }
    )