python之Pyspark 管道中的用户定义转换器

我正在尝试创建一个 pyspark 管道来运行分类模型。我的数据集有一列是一个字符串。因此，在管道中应用模型之前，我使用“StringIndexer”将其转换为数字。

我的管道只包含 2 个阶段 StringIndexer 和 ClassificationModel

StringIndexer 正在创建一个带有索引的新列，但是旧列也被保留。我想在管道中引入一个新的转换器来删除“字符串”列。这可能吗？

有没有其他方法可以删除 StringIndexer 中的实际列？

谢谢

请您参考如下方法:

是的，您可以扩展抽象类 Transformer 并创建您自己的转换器以删除不必要的列。

这应该看起来像下面这样:

import org.apache.spark.ml.Transformer 
import org.apache.spark.ml.param.ParamMap 
import org.apache.spark.ml.util.Identifiable 
import org.apache.spark.sql.{DataFrame, Dataset} 
import org.apache.spark.sql.types.{ 
  ArrayType, 
  StringType, 
  StructField, 
  StructType 
} 
import org.apache.spark.sql.functions.collect_list 
 
class Dropper(override val uid: String) extends Transformer { 
 
  def this() = this(Identifiable.randomUID("dropper")) 
 
  override def transform(dataset: Dataset[_]): DataFrame = { 
    dataset.drop("your-column-name-here") 
  } 
 
  override def copy(extra: ParamMap): Transformer = defaultCopy(extra) 
 
  override def transformSchema(schema: StructType): StructType = { 
    //here you should right your result schema i.e. the schema without the dropped column 
  } 
 
} 

我已经这样做了一段时间，对我来说效果很好。

请注意，您还可以扩展抽象类 Estimator。

希望对您有所帮助。最好的问候