package com.roy.spark;

import org.apache.spark.SparkContext;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.*;
import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.api.java.UDF2;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import scala.collection.Seq;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

import static org.apache.spark.sql.functions.col;

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

public class SparkSQLExample {

    public static void main(String[] args) {
        final SparkSession spark = SparkSession.builder().appName("java SparkSQL Example")
                .master("local[*]")
//                .config(...)
                .getOrCreate();
        final SparkContext sc = spark.sparkContext();
        final Dataset<Row> df = spark.read().json("SparkDemo/src/main/resources/people.json");
        df.show();//查看数据和结构
        df.printSchema();//打印schema消息
        df.select("name").show();//在DataFrame中查看某一列
        //需要的参数是org.apache.spark.sql.Column
        df.filter(col("age").gt(21)).show(); //数据过滤
        df.groupBy(col("age")).count().show();//分组计数。
        //这些逻辑可以看到跟SQL就是很接近的。

        //注册临时表
        df.createOrReplaceTempView("people");
        final Dataset<Row> sqldf = spark.sql("select * from people");
        sqldf.show();

        //在SparkSQL中可以使用UDF 自定义函数。
        //自定义一个只传一个参数的自定义函数。两个泛型，一个是输入参数类型，一个是输出参数类型
        UDF1<Integer,Integer> plusOneFunc = new UDF1<Integer, Integer>() {
            @Override
            public Integer call(Integer x) {
                return x + 1;
            }
        };
        //最后传的这个参数就是函数的输出类型。scala的语法有点怪怪的。
        spark.udf().register("plusOne",plusOneFunc, DataTypes.IntegerType);
        final Dataset<Row> sqldf2 = spark.sql("select plusOne(age) from person");
        sqldf2.show();

        //注册全局临时表
        try {
            df.createGlobalTempView("people");
            spark.sql("SELECT * FROM global_temp.people").show();
            //全局临时表是跨session的。全局有效。
            spark.newSession().sql("SELECT * FROM global_temp.people").show();
        } catch (AnalysisException e) {
            e.printStackTrace();
        }

        //手动创建一个DataSet
        Person person = new Person();
        person.setName("Andy");
        person.setAge(32);
        Encoder<Person> personEncoder = Encoders.bean(Person.class);
        Dataset<Person> javaBeanDS = spark.createDataset(
                Collections.singletonList(person),
                personEncoder
        );
        javaBeanDS.show();

        //创建RDD，从RDD转换成DataFrame
        final JavaRDD<Person> personRDD = spark.read().textFile("SparkDemo/src/main/resources/people.txt").toJavaRDD()
                .map(line -> {
                    final String[] records = line.split(" ");
                    Person person2 = new Person();
                    person2.setName(records[0]);
                    person2.setAge(Integer.parseInt(records[1].trim()));
                    return person2;
                });
        final Dataset<Row> personDF = spark.createDataFrame(personRDD, Person.class);
        personDF.show();

        //DF格式转换
        Encoder<String> stringEncoder = Encoders.STRING();
        Dataset<String> teenagerNamesByIndexDF = personDF.map(
                (MapFunction<Row, String>) row -> "Name: " + row.getString(0),
                stringEncoder);
        teenagerNamesByIndexDF.show();
        //编程式类型转换 无schema的RDD转换成有schema的DF
        JavaRDD<String> personRDD2 = sc
                .textFile("examples/src/main/resources/people.txt", 1)
                .toJavaRDD();
        // The schema is encoded in a string
        String schemaString = "name age";

        // Generate the schema based on the string of schema
        List<StructField> fields = new ArrayList<>();
        for (String fieldName : schemaString.split(" ")) {
            StructField field = DataTypes.createStructField(fieldName, DataTypes.StringType, true);
            fields.add(field);
        }
        StructType schema = DataTypes.createStructType(fields);

        // Convert records of the RDD (people) to Rows
        JavaRDD<Row> rowRDD = personRDD2.map((Function<String, Row>) record -> {
            String[] attributes = record.split(",");
            return RowFactory.create(attributes[0], attributes[1].trim());
        });

        // Apply the schema to the RDD
        Dataset<Row> peopleDataFrame = spark.createDataFrame(rowRDD, schema);
        peopleDataFrame.show();
        //parquet文件缓存DataFrame的数据
//        peopleDataFrame.write().format("parquet").save("peopleDataFrame.parquet");
//        final Dataset<Row> df2 = spark.read().load("peopleDataFrame.parquet");

//        peopleDataFrame.write().parquet("peopleDataFrame.parquet");
//        spark.read().parquet("peopleDataFrame.parquet");
        spark.close();
    }
}
