package org.example;

import org.apache.curator.shaded.com.google.common.hash.BloomFilter;
import org.apache.curator.shaded.com.google.common.hash.Funnels;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.spark.sql.api.java.*;
import org.apache.spark.SparkConf;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.lucene.util.RamUsageEstimator;

/**
 * add by chengwansheng
 */
class MyBloomFilter {
    private BloomFilter bloomFilter;

    public MyBloomFilter(BloomFilter b) {
        bloomFilter = b;
    }

    public BloomFilter getBloomFilter() {
        return bloomFilter;
    }
}

public class BloomUdf implements UDF2<Object, String, Boolean> {
    //最大記錄限制，安全起見
    private static int maxSize = 50000000;

    //布隆過濾器是否開啟配置， 1 開啟，0 關(guān)閉
    private static int udfBloomFilterEnable;

    //布隆過濾器是否開啟參數(shù)，默認(rèn)開啟
    private static String bloomFilterConfKey = "spark.myudf.bloom.enable";

    //加配置配置參數(shù)，目前不起作用?? 
    static {
        SparkConf sparkConf = new SparkConf();
        udfBloomFilterEnable = sparkConf.getInt(bloomFilterConfKey, 1);
        System.out.println("the spark.myudf.bloom.enable value " + udfBloomFilterEnable);
    }

    //布隆過濾器列表，支持多個(gè)布隆過濾器
    private static ConcurrentHashMap<String, MyBloomFilter> bloomFilterMap = new ConcurrentHashMap<>();

    /**
     * 布隆過濾器核心構(gòu)建方法
     * 通過讀取表的 hdfs 文件信息，構(gòu)建布隆過濾器
     * 一個(gè) jvm 只加載一次
     * @param key 
     * @param path 
     * @throws IOException 
     */
    private synchronized static void buildBloomFilter(String key, String path) throws IOException {
        if (!bloomFilterMap.containsKey(key)) {
            BloomFilter bloomFilter;
            Configuration cnotallow=new Configuration();
            FileSystem hdfs=FileSystem.get(conf);
            Path pathDf=new Path(path);
            FileStatus[] stats=hdfs.listStatus(pathDf);

            //獲取記錄總數(shù)
            long sum = 0;
            for (int i=0; i<stats.length; i++){
                InputStream inputStream=hdfs.open(stats[i].getPath());
                InputStreamReader inputStreamReader= new InputStreamReader(inputStream);
                BufferedReader reader=new BufferedReader(inputStreamReader);
                sum = sum + reader.lines().count();
            }

            if(sum > maxSize) {
                //如果數(shù)據(jù)量大于期望值，則將布隆過濾器置空(即布隆過濾器不起作用)
                System.out.println("the max number is " + maxSize + ", but target num is too big, the " + key + " bloom will be invalid");
                bloomFilter = null;
            } else {
                //默認(rèn) 1000 W，超過取樣本數(shù)據(jù) 2 倍的量。這里取 2 倍是為了提高布隆過濾器的效果， 2 倍是一個(gè)比較合適的值
                long exceptSize = sum*2>10000000?sum*2:10000000;
                bloomFilter = BloomFilter.create(Funnels.stringFunnel(StandardCharsets.UTF_8), (int) exceptSize);
                for (int i=0; i<stats.length; i++){
                    //打印每個(gè)文件路徑
                    System.out.println(stats[i].getPath().toString());
                    //讀取每個(gè)文件
                    InputStream inputStream=hdfs.open(stats[i].getPath());
                    InputStreamReader inputStreamReader= new InputStreamReader(inputStream);
                    BufferedReader reader=new BufferedReader(inputStreamReader);
                    String line="";
                    while((line=reader.readLine())!=null){
                        bloomFilter.put(line);
                    }
                }
            }

            MyBloomFilter myBloomFilter = new MyBloomFilter(bloomFilter);
            bloomFilterMap.put(key, myBloomFilter);
            System.out.println("the bloom " + key + " size is " + RamUsageEstimator.humanSizeOf(bloomFilter) + " num " + sum);
        }
    }

    /**
     * 核心調(diào)用方法
     * 參數(shù) s ：被過濾的參數(shù)
     * 參數(shù) key ：需要構(gòu)建的布隆過濾器，此處是庫名 + 表名稱，即 db_name.table_name 
     * @param s 
     * @param key 
     * @return 
     * @throws Exception  
     */
    @Override
    public Boolean call(Object s, String key) throws Exception {
        //如果 spark.myudf.bloom.enable 參數(shù)配置為 0 ，則布隆過濾器失效，直接返回 true 
        if (udfBloomFilterEnable == 0) {
            return true;
        }

        if (!bloomFilterMap.containsKey(key)) {
            String[] table_array = key.split("\\.");
            if (table_array.length != 2) {
                String msg = "the key is invalid: " + key + ", must like db_name.table_name";
                System.out.println(msg);
                throw new IOException(msg);
            }
            String dbName = table_array[0];
            String tableName = table_array[1];
            String path = "/hive/" + dbName + ".db/" + tableName;
            System.out.println(path);
            //構(gòu)建布隆過濾器
            buildBloomFilter(key, path);
        }

        if (!bloomFilterMap.containsKey(key)) {
            String msg = "not found bloom filter " + key;
            System.out.println(msg);
            throw new IOException(msg);
        }

        BloomFilter bloomFilter = bloomFilterMap.get(key).getBloomFilter();
        if (bloomFilter == null) {
            //如果數(shù)據(jù)量大于期望值，則直接返回真，即布隆過濾器不起作用
            return true;
        } else {
            return bloomFilter.mightContain(String.valueOf(s));
        }
    }
}

--建表數(shù)據(jù)
create table default.A (
    item_id bigint comment '商品ID',
    item_name string comment '商品名稱',
    item_price bigint comment '商品價(jià)格',
    create_time timestamp comment '創(chuàng)建時(shí)間',
    update_time timestamp comment '創(chuàng)建時(shí)間'
)

create table default.B (
    item_id bigint comment '商品ID',
    sku_id bigint comment 'skuID',
    sku_price bigint comment '商品價(jià)格',
    create_time timestamp comment '創(chuàng)建時(shí)間',
    update_time timestamp comment '創(chuàng)建時(shí)間'
)

create table default.ot (
    item_id bigint comment '商品ID',
    sku_id bigint comment 'skuID',
    sku_price bigint comment '商品價(jià)格',
    item_price bigint comment '商品價(jià)格'
) PARTITIONED BY (pt string COMMENT '分區(qū)字段') 

--準(zhǔn)備數(shù)據(jù)
insert overwrite table default.A 
values
(1,'測試1',101,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(2,'測試2',102,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(3,'測試2',103,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(4,'測試2',104,'2023-03-25 08:00:00','2023-04-22 08:00:00'),
(5,'測試2',105,'2023-03-25 08:00:00','2023-04-22 08:00:00');

insert overwrite table default.B 
values 
(1,11,201,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(1,12,202,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(1,13,203,'2023-04-22 08:00:00','2023-04-22 08:00:00'),
(2,21,211,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(2,22,212,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(4,42,212,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(5,51,251,'2023-03-25 08:00:00','2023-03-25 08:00:00'),
(5,52,252,'2023-04-22 08:00:00','2023-04-22 08:00:00'),
(5,53,253,'2023-04-22 08:00:00','2023-04-22 08:00:00');

insert overwrite table default.ot partition(pt='20230421')
values 
(1,11,201,101),
(1,12,202,101),
(2,21,211,102),
(2,22,212,102),
(4,42,212,114),
(5,51,251,110);

insert overwrite table default.ot partition(pt='20230422')
select B.item_id 
,B.sku_id 
,B.sku_price 
,A.item_price
from B 
left join A on(A.item_id=B.item_id)

--構(gòu)建布隆過濾器
drop table if exists tmp.tmp_primary_key;
create table tmp.tmp_primary_key stored as TEXTFILE as 
select item_id
from (
    select item_id
    from default.A 
    where update_time>='2023-04-22'
    union all 
    select item_id
    from default.B 
    where update_time>='2023-04-22'
) where length(item_id)>0
group by item_id;

--增量數(shù)據(jù)計(jì)算
insert overwrite table default.ot partition(pt='20230422')
select B.item_id 
,B.sku_id 
,B.sku_price 
,A.item_price
from default.B 
left join default.A on(A.item_id=B.item_id and bloom_filter(A.item_id, "tmp.tmp_primary_key"))
where bloom_filter(B.item_id, "tmp.tmp_primary_key")
union all 
--合并歷史未變更數(shù)據(jù)
select item_id
,sku_id
,sku_price
,item_price
from default.ot
where not bloom_filter(item_id, "tmp.tmp_primary_key")
and pt='20230421'