com.datastax.data.prepare.util.SharedUDFs Maven / Gradle / Ivy

Go to download
package com.datastax.data.prepare.util;

import org.apache.spark.sql.UDFRegistration;
import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.types.DataTypes;

import java.math.BigDecimal;
import java.sql.Date;
import java.sql.Timestamp;
import java.text.DecimalFormat;
import java.util.UUID;

public class SharedUDFs {

    public static void uuid(UDFRegistration udfRegistration) {
//        问题：java.lang.IllegalArgumentException: Invalid lambda deserialization
//              at com.datastax.data.prepare.util.c.$deserializeLambda$(SharedUDFs.java:14)
//        解决：https://sourceforge.net/p/proguard/bugs/545/
//        udfRegistration.register("uuid", (Object obj) -> UUID.randomUUID().toString().replaceAll("-", ""), DataTypes.StringType);  //

        udfRegistration.register("uuid", new UDF1() {
            @Override
            public String call(Object a) throws Exception {
                return UUID.randomUUID().toString().replaceAll("-", "");
            }
        }, DataTypes.StringType);
    }

    /**
     * num2date, 列为Long型
     * @param udfRegistration  注册udf
     */
    public static void num2date(UDFRegistration udfRegistration) {
        udfRegistration.register("num2date", new UDF1() {
            @Override
            public Timestamp call(Long a) throws Exception {
                return new Timestamp(new java.util.Date(a).getTime());
            }
        }, DataTypes.TimestampType);
    }

    /**
     * timestamp2num, 列为Timestamp型
     * @param udfRegistration  注册udf
     */
    public static void timestamp2num(UDFRegistration udfRegistration) {
        udfRegistration.register("timestamp2num", new UDF1() {
            @Override
            public Long call(Timestamp timestamp) throws Exception {
                return timestamp.getTime();
            }
        }, DataTypes.LongType);
    }

    /**
     * date2num, 列为Date型
     * @param udfRegistration  注册udf
     */
    public static void date2num(UDFRegistration udfRegistration) {
        udfRegistration.register("date2num", new UDF1() {
            @Override
            public Long call(Date date) throws Exception {
                return date.getTime();
            }
        }, DataTypes.LongType);
    }

    /**
     * 按照规则取舍, 列为数值型
     * @param udfRegistration  注册udf
     * @param numberFormat  形式为 ###.00 等
     * @param roundingMode  BigDecimal 的 RoundingMode
     */
    public static void rounding(UDFRegistration udfRegistration, String numberFormat, String roundingMode) {
        DecimalFormat decimalFormat = new DecimalFormat(numberFormat);
        udfRegistration.register("rounding", new UDF1() {
            @Override
            public String call(Object object) throws Exception {
                BigDecimal bigDecimal = new BigDecimal(object.toString());
                bigDecimal.setScale(numberFormat.length() - numberFormat.lastIndexOf(".") - 1, SharedMethods.getRoundMode(roundingMode));
                return decimalFormat.format(bigDecimal);
            }
        }, DataTypes.StringType);
    }

    /**
     * string2dateString, 不够完善，以后可以支持类似正则匹配的方式转换成date //todo
     * @param udfRegistration 注册udf
     */
    public static void string2date(UDFRegistration udfRegistration) {
        udfRegistration.register("string2dateString", new UDF1() {
            @Override
            public String call(String s) throws Exception {
                s = s.trim();
                if(s.length() == 0) {
                    return s;
                }
                if(s.matches("\\d+\\.?\\d+")) {
                    int position = s.indexOf(".");
                    if(position == -1) {
                        position = s.length();
                    }
                    long l = Long.parseLong(s.substring(0, position));
                    return new Timestamp(l).toString();
                }
                return s;
            }
        }, DataTypes.StringType);
    }

}