• Example Search
• Project Search

• Java
• Python
• JavaScript
• TypeScript
• C++
• Scala
• Blog

• aardpfark-master
  • src
    • main
      • scala
        • sparktypes.scala
        • com
          • ibm
            • aardpfark
              • avro
                • SchemaConverters.scala
              • spark
                • ml
                  • classification
                    • DecisionTreeClassificationModel.scala
                    • MLPClassifier.scala
                    • RandomForestClassificationModel.scala
                    • LogisticRegressionModel.scala
                    • NaiveBayesModel.scala
                    • LinearSVCModel.scala
                    • GBTClassificationModel.scala
                  • PFATransformers.scala
                  • clustering
                    • KMeans.scala
                  • linear
                    • LinearModelData.scala
                    • SparkLinearModel.scala
                  • Merge.scala
                  • tree
                    • Trees.scala
                  • regression
                    • GBTRegressionModel.scala
                    • RandomForestRegressionModel.scala
                    • GeneralizedLinearRegressionModel.scala
                    • LinearRegressionModel.scala
                    • DecisionTreeRegressionModel.scala
                  • feature
                    • Binarizer.scala
                    • StringIndexerModel.scala
                    • StopWordsRemover.scala
                    • MaxAbsScaler.scala
                    • Bucketizer.scala
                    • ElementwiseProduct.scala
                    • MinMaxScaler.scala
                    • VectorSelector.scala
                    • PCAModel.scala
                    • Normalizer.scala
                    • CountVectorizerModel.scala
                    • RegexTokenizer.scala
                    • ChiSqSelectorModel.scala
                    • VectorAssembler.scala
                    • NGram.scala
                    • VectorSlicer.scala
                    • StandardScaler.scala
                    • IDF.scala
                  • SparkSupport.scala
              • pfa
                • functions
                  • package.scala
                  • FunctionLibrary.scala
                • utils
                  • Util.scala
                • package.scala
                • document
                  • JSONSerializers.scala
                  • PFABuilder.scala
                  • Cell.scala
                  • PFADocument.scala
                • types
                  • package.scala
                  • WithSchema.scala
                • expression
                  • LetSet.scala
                  • FunctionRefs.scala
                  • ControlStructures.scala
                  • Casts.scala
                  • package.scala
                  • PFAExpression.scala
                  • FunctionCalls.scala
                  • New.scala
                  • AttributeRetrieval.scala
                  • Loops.scala
    • test
      • scala
        • com
          • ibm
            • aardpfark
              • spark
                • ml
                  • classification
                    • DecisionTreeClassifierSuite.scala
                    • MLPClassifierSuite.scala
                    • SparkClassifierPFASuiteBase.scala
                    • GBTClassificationSuite.scala
                    • RandomForestClassificationSuite.scala
                    • NaiveBayesSuite.scala
                    • LinearSVCSuite.scala
                    • LogisticRegressionSuite.scala
                  • SparkMLTestUtils.scala
                  • clustering
                    • KMeansModelSuite.scala
                  • PredictorPipelineSuite.scala
                  • regression
                    • SparkRegressorPFASuiteBase.scala
                    • LinearRegressionSuite.scala
                    • DecisionTreeRegressionSuite.scala
                    • GeneralizedLinearRegressionSuite.scala
                    • RandomForestRegressionSuite.scala
                    • GBTRegressionSuite.scala
                  • feature
                    • MinMaxScalerSuite.scala
                    • NormalizerSuite.scala
                    • NGramSuite.scala
                    • StopWordsRemoverSuite.scala
                    • VectorAssemblerSuite.scala
                    • ElementwiseProductSuite.scala
                    • VectorSlicerSuite.scala
                    • RegexTokenizerSuite.scala
                    • PCAModelSuite.scala
                    • ChiSqSelectorSuite.scala
                    • IDFSuite.scala
                    • StringIndexerModelSuite.scala
                    • StandardScalerSuite.scala
                    • CountVectorizerSuite.scala
                    • MaxAbsScalerSuite.scala
                    • BucketizerSuite.scala
                    • BinarizerSuite.scala
              • pfa
                • functions
                  • StringLibrarySuite.scala
                  • ImputeLibrarySuite.scala
                  • MapLibrarySuite.scala
                  • LinearAlgebraLibrarySuite.scala
                  • MathLibrarySuite.scala
                  • RegexLibrarySuite.scala
                  • CastLibrarySuite.scala
                  • FunctionLibrarySuite.scala
                  • ArrayLibrarySuite.scala
                  • CoreLibrarySuite.scala
                • SparkPFASuiteBase.scala
                • SparkFeaturePFASuiteBase.scala
                • SparkPredictorPFASuiteBase.scala
                • DSLSuiteBase.scala
                • PFATestUtils.scala
                • expression
                  • FunctionSuite.scala
                  • CastsSuite.scala
                  • LetSetSuite.scala
                  • AttrSuite.scala
                  • ControlStructuresSuite.scala
                  • LoopsSuite.scala
  • MAINTAINERS.md
  • LICENSE
  • project
    • build.properties
    • plugins.sbt
  • CONTRIBUTING.md
  • python
    • README.md
  • data
    • sample_libsvm_data.txt
    • sample_multiclass_classification_data.txt
    • sample_kmeans_data.txt
    • sample_binary_classification_data.txt
    • sample_linear_regression_data.txt
    • sample_lda_libsvm_data.txt
  • .travis.yml
  • README.md
  • build.sbt
  • .gitignore

package com.ibm.aardpfark.spark.ml.feature

import com.ibm.aardpfark.pfa.{Result, SparkFeaturePFASuiteBase}
import org.apache.spark.ml.feature.CountVectorizer

class CountVectorizerSuite extends SparkFeaturePFASuiteBase[CountVectorizerResult] {

  val df = spark.createDataFrame(Seq(
    (0, Array("a", "b", "c", "d", "e", "f")),
    (1, Array("a", "b", "b", "c", "a"))
  )).toDF("id", "words")

  val cv = new CountVectorizer()
    .setInputCol("words")
    .setOutputCol("features")

  override val sparkTransformer = cv.fit(df)

  val result = sparkTransformer.transform(df)
  override val input = result.select(cv.getInputCol).toJSON.collect()
  override val expectedOutput = withColumnAsArray(result, cv.getOutputCol).toJSON.collect()

  // Additional test for MinTF
  test("CountVectorizer with MinTF = 0.3") {
    val cv = new CountVectorizer()
      .setInputCol("words")
      .setOutputCol("features")
      .setMinTF(0.3)

    val sparkTransformer = cv.fit(df)

    val result = sparkTransformer.transform(df)
    val input = result.select(cv.getInputCol).toJSON.collect()
    val expectedOutput = withColumnAsArray(result, cv.getOutputCol).toJSON.collect()

    parityTest(sparkTransformer, input, expectedOutput)
  }

  test("CountVectorizer with MinTF = 2.0") {
    val cv = new CountVectorizer()
      .setInputCol("words")
      .setOutputCol("features")
      .setMinTF(2.0)

    val sparkTransformer = cv.fit(df)

    val result = sparkTransformer.transform(df)
    val input = result.select(cv.getInputCol).toJSON.collect()
    val expectedOutput = withColumnAsArray(result, cv.getOutputCol).toJSON.collect()

    parityTest(sparkTransformer, input, expectedOutput)
  }

  // Additional test for binary
  test("CountVectorizer with binary") {
    val cv = new CountVectorizer()
      .setInputCol("words")
      .setOutputCol("features")
      .setBinary(true)

    val sparkTransformer = cv.fit(df)

    val result = sparkTransformer.transform(df)
    val input = result.select(cv.getInputCol).toJSON.collect()
    val expectedOutput = withColumnAsArray(result, cv.getOutputCol).toJSON.collect()

    parityTest(sparkTransformer, input, expectedOutput)
  }

}

case class CountVectorizerResult(features: Seq[Double]) extends Result