Сравнение значения столбца с целым числом (переменная итерации цикла) при использовании условия when в withColumn

Question

При использовании withColumn для добавления столбца в фрейм данных мне нужно сравнить значение столбца с переменной итерации цикла, чтобы присвоить значение новому столбцу.Даже если мои условия для первого или второго момента могут быть выполнены, цикл всегда выполняет другое условие.

Я пытался использовать lit (j)

j), j

for(j <- 1 to 120) {

  actuarialEvaluationDF = actuarialEvaluationDF.withColumn("PolAnn(" + j + ")",
    when(lit(j) === lit(1), monthAdd(col("HOISSDTE"), lit(1)))
      .when(col(completeMonth) > 0 && j > 0 && expr("completeMonth > " + j) , monthAdd(col("HOISSDTE"), lit(j))) // CompleteMonth > 0 and CompleteMonth > j > 0
      .otherwise(col("PolAnn(" + (j-1) + ")" ))) //TODO: Confirm
}

Ожидаемый результат состоит в том, что значение monthAdd (col ("HOISSDTE"), lit (j))), когда j больше 0 и меньше значения столбца в completeMonth.Когда j> значение столбца completeMonth, необходимо вернуть значение Pol (Ann (j-1)).