Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Appends prediction columns to transform schema #60

Merged
merged 2 commits into from
Dec 17, 2024
Merged

Appends prediction columns to transform schema #60

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f61607b
Appends prediction columns to transform schema
Marcus-Rosti Nov 27, 2024
4d49dca
fixes the comment
Marcus-Rosti Nov 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions isolation-forest/src/main/scala/com/linkedin/relevance/isolationforest/IsolationForest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import org.apache.spark.ml.linalg.Vector
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, Identifiable}
import org.apache.spark.ml.Estimator
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
import org.apache.spark.sql.Dataset
import org.apache.spark.{HashPartitioner, TaskContext}

Expand Down Expand Up @@ -187,8 +187,8 @@ class IsolationForest(override val uid: String) extends Estimator[IsolationFores

/**
* Validates the input schema and transforms it into the output schema. It validates that the
* input DataFrame has a $(featuresCol) of the correct type. In this case, the output schema is
* identical to the input schema.
* input DataFrame has a $(featuresCol) of the correct type. In this case, the output schema appends
* the output columns to the input schema.
*
* @param schema The schema of the DataFrame containing the data to be fit.
* @return The schema of the DataFrame containing the data to be fit.
Expand All @@ -200,7 +200,16 @@ class IsolationForest(override val uid: String) extends Estimator[IsolationFores
require(schema($(featuresCol)).dataType == VectorType,
s"Input column ${$(featuresCol)} is not of required type ${VectorType}")

val outputFields = schema.fields
val outputFields: Array[StructField] = schema.fields ++ Array(
StructField(
name = s"$predictionCol",
dataType = DoubleType
),
StructField(
name = s"$scoreCol",
dataType = DoubleType
)
)

StructType(outputFields)
}
Expand Down
Loading