com.databricks.labs.automl.exploration.structures.Utilities.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of automatedml_2.11 Show documentation

Databricks Labs AutoML toolkit

The newest version!

package com.databricks.labs.automl.exploration.structures

import org.apache.spark.sql.DataFrame

object CutoffTypes extends Enumeration {
  type CutoffTypes = Value
  val None, Threshold, Count = Value
}

object FeatureImportanceModelFamily extends Enumeration {
  type FeatureImportanceModelFamily = Value
  val RandomForest, XGBoost = Value
}

object ModelType extends Enumeration {
  type ModelType = Value
  val Regressor, Classifier = Value
}

sealed trait FIReturn {
  def data: DataFrame
  def fieldsInVector: Array[String]
  def allFields: Array[String]
}

abstract case class FeatureImportanceOutput() extends FIReturn
abstract case class FeatureImportanceReturn(importances: DataFrame,
                                            topFields: Array[String])
    extends FIReturn