package com.roy.spark.ml;

import org.apache.spark.ml.feature.CountVectorizer;
import org.apache.spark.ml.feature.CountVectorizerModel;
import org.apache.spark.mllib.linalg.Matrices;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.*;

import java.util.Arrays;
import java.util.List;

/**
 * @author ：楼兰
 * @date ：Created in 2021/5/25
 * @description:
 **/

public class CountVecDemo {
    public static void main(String[] args) {
        SparkSession spark = SparkSession
                .builder()
                .master("local")
                .appName("CountVecDemo")
                .getOrCreate();

        List<Row> data = Arrays.asList(
                RowFactory.create(Arrays.asList("love", "java", "love","tl")),
                RowFactory.create(Arrays.asList("java", "is", "the", "best")),
                RowFactory.create(Arrays.asList("tl", "is", "the", "best","java","tech"))
        );
        StructType schema = new StructType(new StructField[] {
                new StructField("text", new ArrayType(DataTypes.StringType, true), false, Metadata.empty())
        });
        Dataset<Row> df = spark.createDataFrame(data, schema);

        // fit a CountVectorizerModel from the corpus
        CountVectorizerModel cvModel = new CountVectorizer()
                .setInputCol("text")
                .setOutputCol("feature")
                .setVocabSize(3)
                .setMinDF(2)
                .fit(df);

        // alternatively, define CountVectorizerModel with a-priori vocabulary
        CountVectorizerModel cvm = new CountVectorizerModel(new String[]{"java","love","tl","the"})
                .setInputCol("text")
                .setOutputCol("feature");

        final Dataset<Row> resDf = cvModel.transform(df);
        resDf.show(false);
        // $example off$

        spark.stop();
    }
}
