Skip to content
Snippets Groups Projects
Commit 111a6247 authored by Yanbo Liang's avatar Yanbo Liang Committed by Xiangrui Meng
Browse files

[SPARK-14147][ML][SPARKR] SparkR predict should not output feature column

## What changes were proposed in this pull request?
SparkR does not support type of vector which is the default type of feature column in ML. R predict also does not output intermediate feature column. So SparkR ```predict``` should not output feature column. In this PR, I only fix this issue for ```naiveBayes``` and ```survreg```. ```kmeans``` has the right code route already and  ```glm``` will be fixed at SparkRWrapper refactor(#12294).

## How was this patch tested?
No new tests.

cc mengxr shivaram

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #11958 from yanboliang/spark-14147.
parent 1995c2e6
No related branches found
No related tags found
No related merge requests found
......@@ -44,7 +44,7 @@ private[r] class AFTSurvivalRegressionWrapper private (
}
def transform(dataset: Dataset[_]): DataFrame = {
pipeline.transform(dataset)
pipeline.transform(dataset).drop(aftModel.getFeaturesCol)
}
}
......
......@@ -37,7 +37,9 @@ private[r] class NaiveBayesWrapper private (
lazy val tables: Array[Double] = naiveBayesModel.theta.toArray.map(math.exp)
def transform(dataset: Dataset[_]): DataFrame = {
pipeline.transform(dataset).drop(PREDICTED_LABEL_INDEX_COL)
pipeline.transform(dataset)
.drop(PREDICTED_LABEL_INDEX_COL)
.drop(naiveBayesModel.getFeaturesCol)
}
}
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment