viirya commented on a change in pull request #26547: [SPARK-29914][ML] ML
models attach metadata in `transform`/`transformSchema`
URL: https://github.com/apache/spark/pull/26547#discussion_r349906824
##########
File path: mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala
##########
@@ -106,6 +107,93 @@ private[spark] object SchemaUtils {
StructType(schema.fields :+ col)
}
+ /**
+ * Update the size of a ML Vector column. If this column do not exist,
append it.
+ * @param schema input schema
+ * @param colName column name
+ * @param size number of features
+ * @return new schema
+ */
+ def updateAttributeGroupSize(
+ schema: StructType,
+ colName: String,
+ size: Int): StructType = {
+ require(size > 0)
+ val attrGroup = new AttributeGroup(colName, size)
+ val field = attrGroup.toStructField
+ updateField(schema, field, true)
+ }
+
+ /**
+ * Update the number of values of an existing column. If this column do not
exist, append it.
+ * @param schema input schema
+ * @param colName column name
+ * @param numValues number of values.
+ * @return new schema
+ */
+ def updateNumValues(
+ schema: StructType,
+ colName: String,
+ numValues: Int): StructType = {
+ val attr = if (numValues == 2) {
+ BinaryAttribute.defaultAttr
+ .withName(colName)
+ } else {
+ NominalAttribute.defaultAttr
+ .withName(colName)
+ .withNumValues(numValues)
+ }
Review comment:
Not sure about this. Is numValues == 2 case always BinaryAttribute?
NominalAttribute can not have two number of values?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
