Spark: работа с нулем в функции преобразования tfidf

Question

Я пытаюсь сохранить информацию о значении null в столбце после применения преобразования (преобразования HashingTF и IDFModel) столбца данных.

Таким образом, регистр df.na.drop(column_name) не должен использоваться до преобразования.

Выполнение следующих действий дает Caused by: java.lang.NullPointerException, как и ожидалось:

import org.apache.spark.ml.feature.{HashingTF, IDFModel}
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions.{col, lit}

def seq2tfidf(df: DataFrame, tf: HashingTF, idf: IDFModel, column: String): DataFrame = {
  val renamedDf = df.withColumnRenamed(column, "value")
  val tfValues = tf.transform(withValueDf)
  val idfValues = idf.transform(tfValues)
    .withColumnRenamed("idfFeature", column)
    .drop("tfFeature")
    .drop("value")

  idfValues
}

То, что работает, это разделение на нуль и ноль, поэтому в 2 кадра данных и затем объединение в конце. Как следующее:

def seq2tfidf(df: DataFrame, tf: HashingTF, idf: IDFModel, column: String): DataFrame = {
  val renamedDf = df.withColumnRenamed(column, "value")

  val withValueDf = renamedDf.filter(col("value").isNotNull)
  val withNullValueDf = renamedDf
    .filter(col("value").isNull)
    .withColumn("idfFeature", lit(null: SparseVector))
    .drop("value")

  val tfValues = tf.transform(withValueDf)
  val idfValues = idf.transform(tfValues)
    .withColumnRenamed("idfFeature", column)
    .drop("tfFeature")
    .drop("value")

  idfValues.union(withNullValueDf)
}

Работает, как ожидалось, но, кажется, немного медленно. Есть ли другой способ сделать это?