<commit_before>package com.aluxian.tweeather.transformers

import org.apache.spark.ml.Transformer
import org.apache.spark.ml.param._
import org.apache.spark.ml.util.{BasicParamsReadable, BasicParamsWritable, Identifiable}
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.types._

/**
  * A transformer that removes columns.
  */
class ColumnDropper(override val uid: String) extends Transformer with BasicParamsWritable {

  def this() = this(Identifiable.randomUID("columnsDropper"))

  /**
    * Param for the column names to be removed.
    * @group param
    */
  final val dropColumns: StringArrayParam =
    new StringArrayParam(this, "dropColumns", "columns to be dropped")

  /** @group setParam */
  def setDropColumns(columns: String*): this.type = set(dropColumns, columns.toArray)

  /** @group setParam */
  def setDropColumns(columns: Array[String]): this.type = set(dropColumns, columns)

  /** @group getParam */
  def getDropColumns: Array[String] = $(dropColumns)

  setDefault(dropColumns -> Array())

  override def transformSchema(schema: StructType): StructType = {
    StructType(schema.fields.diff($(dropColumns)))
  }

  override def transform(dataset: DataFrame): DataFrame = {
    transformSchema(dataset.schema, logging = true)
    dataset.drop($(dropColumns): _*)
  }

  override def copy(extra: ParamMap): ColumnDropper = defaultCopy(extra)

}

object ColumnDropper extends BasicParamsReadable[ColumnDropper] {
  override def load(path: String): ColumnDropper = super.load(path)
}
<commit_msg>Fixes usage of df.drop in column dropper<commit_after>package com.aluxian.tweeather.transformers

import com.aluxian.tweeather.RichArray
import org.apache.spark.ml.Transformer
import org.apache.spark.ml.param._
import org.apache.spark.ml.util.{BasicParamsReadable, BasicParamsWritable, Identifiable}
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.types._

/**
  * A transformer that removes columns.
  */
class ColumnDropper(override val uid: String) extends Transformer with BasicParamsWritable {

  def this() = this(Identifiable.randomUID("columnsDropper"))

  /**
    * Param for the column names to be removed.
    * @group param
    */
  final val dropColumns: StringArrayParam =
    new StringArrayParam(this, "dropColumns", "columns to be dropped")

  /** @group setParam */
  def setDropColumns(columns: String*): this.type = set(dropColumns, columns.toArray)

  /** @group setParam */
  def setDropColumns(columns: Array[String]): this.type = set(dropColumns, columns)

  /** @group getParam */
  def getDropColumns: Array[String] = $(dropColumns)

  setDefault(dropColumns -> Array())

  override def transformSchema(schema: StructType): StructType = {
    StructType(schema.fields.diff($(dropColumns)))
  }

  override def transform(dataset: DataFrame): DataFrame = {
    transformSchema(dataset.schema, logging = true)
    $(dropColumns).mapCompose(dataset)(col => _.drop(col))
  }

  override def copy(extra: ParamMap): ColumnDropper = defaultCopy(extra)

}

object ColumnDropper extends BasicParamsReadable[ColumnDropper] {
  override def load(path: String): ColumnDropper = super.load(path)
}

{'fraction_non_alphanumeric': 0.11773989898989899, 'fraction_numerical': 0.0, 'mean_word_length': 4.814678899082569, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '419344', 'n_tokens_mistral': 957, 'n_tokens_neox': 938, 'n_words': 223}