PredictLinks.scala

// probability will be a vector with 2 values resp. for label 0.0 and 1.0. We just need the 1.0 one
val getP1 = udf((v: Vector) => v(1))

val predictionsRaw = lrModel.transform(comparableDataset)

val predictedLinksDataset = predictionsRaw
      .select(
        $"left.dossierNummer".as("dossierNummer_left"),
        $"right.dossierNummer".as("dossierNummer_right"),
        $"features",
        getP1($"probability").as("probability"),
        $"prediction".as("label")
      )
      .filter('label === 1.0)
      .as[PredictedVector]