Skip to content
Snippets Groups Projects
Commit 47c1d562 authored by Mike Dusenberry's avatar Mike Dusenberry Committed by Joseph K. Bradley
Browse files

[SPARK-7426] [MLLIB] [ML] Updated Attribute.fromStructField to allow any NumericType.

Updated `Attribute.fromStructField` to allow any `NumericType`, rather than just `DoubleType`, and added unit tests for a few of the other NumericTypes.

Author: Mike Dusenberry <dusenberrymw@gmail.com>

Closes #6540 from dusenberrymw/SPARK-7426_AttributeFactory.fromStructField_Should_Allow_NumericTypes and squashes the following commits:

87fecb3 [Mike Dusenberry] Updated Attribute.fromStructField to allow any NumericType, rather than just DoubleType, and added unit tests for a few of the other NumericTypes.
parent a1894422
No related branches found
No related tags found
No related merge requests found
......@@ -20,7 +20,7 @@ package org.apache.spark.ml.attribute
import scala.annotation.varargs
import org.apache.spark.annotation.DeveloperApi
import org.apache.spark.sql.types.{DoubleType, Metadata, MetadataBuilder, StructField}
import org.apache.spark.sql.types.{DoubleType, NumericType, Metadata, MetadataBuilder, StructField}
/**
* :: DeveloperApi ::
......@@ -127,7 +127,7 @@ private[attribute] trait AttributeFactory {
* Creates an [[Attribute]] from a [[StructField]] instance.
*/
def fromStructField(field: StructField): Attribute = {
require(field.dataType == DoubleType)
require(field.dataType.isInstanceOf[NumericType])
val metadata = field.metadata
val mlAttr = AttributeKeys.ML_ATTR
if (metadata.contains(mlAttr)) {
......
......@@ -215,5 +215,10 @@ class AttributeSuite extends SparkFunSuite {
assert(Attribute.fromStructField(fldWithoutMeta) == UnresolvedAttribute)
val fldWithMeta = new StructField("x", DoubleType, false, metadata)
assert(Attribute.fromStructField(fldWithMeta).isNumeric)
// Attribute.fromStructField should accept any NumericType, not just DoubleType
val longFldWithMeta = new StructField("x", LongType, false, metadata)
assert(Attribute.fromStructField(longFldWithMeta).isNumeric)
val decimalFldWithMeta = new StructField("x", DecimalType(None), false, metadata)
assert(Attribute.fromStructField(decimalFldWithMeta).isNumeric)
}
}
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