ai.catboost.spark.impl.Helpers.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of catboost-spark_2.4_2.12 Show documentation

JVM module to use CatBoost on Apache Spark

The newest version!

package ai.catboost.spark.impl

import java.nio.file.{Files,Path}
import java.nio.charset.StandardCharsets
import java.util.concurrent.{ExecutorCompletionService,Future}


private[spark] object Helpers {
  def writeToTempFile(
    content: String,
    tmpFilePrefix: String,
    tmpFileSuffix: Option[String] = None
  ) : Path = {
    val tempFile = Files.createTempFile(tmpFilePrefix, tmpFileSuffix.getOrElse(".tmp"))
    tempFile.toFile.deleteOnExit
    Files.write(tempFile, content.getBytes(StandardCharsets.UTF_8))
    tempFile
  }

  def checkOneFutureAndWaitForOther(
    future1: Future[Unit],
    future1Description: String,
    future2: Future[Unit],
    future2Description: String
  ) = {
    try {
      future1.get
    } catch {
      case e : java.util.concurrent.ExecutionException => {
        future2.cancel(true)
        throw new java.util.concurrent.ExecutionException(
          "Error while executing " + future1Description,
          e.getCause
        )
      }
    }
    try {
      future2.get
    } catch {
      case e : java.util.concurrent.ExecutionException => {
        throw new java.util.concurrent.ExecutionException(
          "Error while executing " + future2Description,
          e.getCause
        )
      }
    }
  }

  def waitForTwoFutures(
    ecs: ExecutorCompletionService[Unit],
    future1: Future[Unit],
    future1Description: String,
    future2: Future[Unit],
    future2Description: String
  ) = {
    val firstCompletedFuture = ecs.take()

    if (firstCompletedFuture == future1) {
      checkOneFutureAndWaitForOther(future1, future1Description, future2, future2Description)
    } else {
      checkOneFutureAndWaitForOther(future2, future2Description, future1, future1Description)
    }
  }
}