package com.roy.spark.ml;

import org.apache.spark.ml.feature.OneHotEncoder;
import org.apache.spark.ml.feature.OneHotEncoderModel;
import org.apache.spark.ml.feature.StringIndexer;
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.DataTypes;
import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

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

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

public class OneHotDemo {
    public static void main(String[] args) {
        SparkSession spark = SparkSession
                .builder()
                .master("local[*]")
                .appName("JavaOneHotEncoderExample")
                .getOrCreate();
        // Note: categorical features are usually first encoded with StringIndexer
        // $example on$
        List<Row> data = Arrays.asList(
                RowFactory.create("北京", 28),
                RowFactory.create("上海", 30),
                RowFactory.create("天津", 29),
                RowFactory.create("长沙", 37),
                RowFactory.create("武汉", 40),
                RowFactory.create("深圳", 35)
        );


        StructType schema = new StructType(new StructField[]{
                new StructField("city", DataTypes.StringType, false, Metadata.empty()),
                new StructField("temperature", DataTypes.IntegerType, false, Metadata.empty())
        });

        Dataset<Row> df = spark.createDataFrame(data, schema);

        final StringIndexer indexer = new StringIndexer().setInputCol("city").setOutputCol("cityIndex");
        final Dataset<Row> df_indexed = indexer.fit(df).transform(df);
        System.out.println("============== step 1 ===============");
        df_indexed.show();

        OneHotEncoder encoder = new OneHotEncoder()
                .setInputCols(new String[] {"cityIndex","temperature"})
                .setOutputCols(new String[] { "cityVec","tempVec"});

        OneHotEncoderModel model = encoder.setDropLast(false).fit(df_indexed);
        Dataset<Row> encoded = model.transform(df_indexed);
        encoded.show();
        // $example off$
        /*
         +----+-----------+---------+-------------+---------------+
        |city|temperature|cityIndex|      cityVec|        tempVec|
        +----+-----------+---------+-------------+---------------+
        |北京|         28|      1.0|(6,[1],[1.0])|(41,[28],[1.0])|
        |上海|         30|      0.0|(6,[0],[1.0])|(41,[30],[1.0])|
        |天津|         29|      2.0|(6,[2],[1.0])|(41,[29],[1.0])|
        |长沙|         37|      5.0|(6,[5],[1.0])|(41,[37],[1.0])|
        |武汉|         40|      3.0|(6,[3],[1.0])|(41,[40],[1.0])|
        |深圳|         35|      4.0|(6,[4],[1.0])|(41,[35],[1.0])|
        cityVec 字典值，构建了6*6的矩阵
         (6,[1],[1.0]) 向量长度为6，索引1的位置，值是1.0
         tempVec 数字型，构建了6*41的矩阵
         (41,[28],[1.0]) 向量长度41，索引为28的位置，值是1.0
         */

        spark.stop();
    }
}
