2021-05-26

概念及核心

Spark Catalyst扩展点

Spark catalyst的扩展点在SPARK-18127中被引入，Spark用户可以在SQL处理的各个阶段扩展自定义实现，非常强大高效，下面我们具体看看其提供的接口和在Spark中的实现。

SparkSessionExtensions

SparkSessionExtensions保存了所有用户自定义的扩展规则，自定义规则保存在成员变量中，对于不同阶段的自定义规则，SparkSessionExtensions提供了不同的接口。

新增自定义规则

用户可以通过SparkSessionExtensions提供的inject开头的方法添加新的自定义规则，具体的inject接口如下：

• injectOptimizerRule – 添加optimizer自定义规则，optimizer负责逻辑执行计划的优化。
• injectParser – 添加parser自定义规则，parser负责SQL解析。
• injectPlannerStrategy – 添加planner strategy自定义规则，planner负责物理执行计划的生成。
• injectResolutionRule – 添加Analyzer自定义规则到Resolution阶段，analyzer负责逻辑执行计划生成。
• injectPostHocResolutionRule – 添加Analyzer自定义规则到Post Resolution阶段。
• injectCheckRule – 添加Analyzer自定义Check规则。

配置自定义规则

在Spark中，用户自定义的规则可以通过两种方式配置生效：

1. 使用SparkSession.Builder中的withExtenstion方法，withExtension方法是一个高阶函数，接收一个自定义函数作为参数，这个自定义函数以SparkSessionExtensions作为参数，用户可以实现这个函数，通过SparkSessionExtensions的inject开头的方法添加用户自定义规则。
2. 通过Spark配置参数，具体参数名为spark.sql.extensions。用户可以将1中的自定义函数实现定义为一个类，将完整类名作为参数值。

代码片

//入口
public void sqlParserV2 (String sparkSql) {

        SparkSession sparkSession = getSparkSession();
        ParserInterface sqlParser = sparkSession.sessionState().sqlParser();
        try {
            sqlParser.parsePlan(sparkSql);
        } catch (ParseException parseException) {
            throw new LyraException("sql语法错误，请输入正确的sql语句");
        } catch (LyraException lyraEx) {
            throw new LyraException(lyraEx.getMessage());
        } catch (Exception e) {
            throw new LyraException("校验sql异常,请联系管理人员");
        }finally {
            sparkSession.stop();
        }

private SparkSession getSparkSession () {
     return SparkSession.builder()
             .withExtensions(new FunctionDo())
             .master("local")
             .appName("test")
             .config("spark.testing.memory","2147480000")
             .getOrCreate();
 }

class FunctionDo implements Function1<SparkSessionExtensions, BoxedUnit> {

     @Override
     public BoxedUnit apply(SparkSessionExtensions sparkSessionExtensions) {
         sparkSessionExtensions.injectParser(new Function2Do());
         return BoxedUnit.UNIT;
     }
  	//..........
 }     

class Function2Do implements Function2<SparkSession, ParserInterface, ParserInterface> {

       @Override
       public StrictParser apply(SparkSession sparkSession, ParserInterface parserInterface) {

           return  new StrictParser(parserInterface);
       }
   	//..........
   }

class StrictParser implements ParserInterface {

        private ParserInterface parserInterface;

        public StrictParser (ParserInterface parserInterface) {
            this.parserInterface = parserInterface;
        }

        private LogicalPlan do_parsePlan_to_find_project(LogicalPlan logicalPlan) {
            if (logicalPlan instanceof Project) {
                return logicalPlan;
            }
            return do_parsePlan_to_find_project(logicalPlan.children().head());
        }

        @Override
        public LogicalPlan parsePlan(String s) throws ParseException {
            LogicalPlan logicalPlan = parserInterface.parsePlan(s);

            logicalPlan = do_parsePlan_to_find_project(logicalPlan);
            String s1 = logicalPlan.toJSON();
            List<String> columns = new ArrayList<>();
            JSONArray array = JSONArray.parseArray(s1);
            JSONObject obj = array.getJSONObject(0);
            JSONArray columnArr = obj.getJSONArray("projectList");
            for (int i = 0; i < columnArr.size(); i++) {
                JSONArray columnArr2 = columnArr.getJSONArray(i);
                int count = 0;
                for (int j = 0; j < columnArr2.size(); j++) {
                    JSONObject obj2 = columnArr2.getJSONObject(j);
                    if ("org.apache.spark.sql.catalyst.expressions.Alias".equals(obj2.getString("class"))) {
                        String column = obj2.getString("name");
                        columns.add(column);
                        count ++;
                        break;
                    }
                }
                if (0 == count) {
                    String column = columnArr2.getJSONObject(0).getString("nameParts");
                    columns.add(column.substring(1,column.length() -1));
                }
            }
            boolean exitDis = columns.stream().anyMatch(z -> z.equals("distinct_id"));
            boolean exitValue = columns.stream().anyMatch(z -> z.equals("value"));
            if (!exitDis || !exitValue) {
                throw new LyraException("SELECT语句必须限定为 distinct_id 和 value 两个投影字段");
            }
            return logicalPlan;
        }

        //..............
    }  

上面的例子为校验sql的查询字段必须包含某几个字段或某几个字段的映射

*[注] 本文提供Java方式实现思路，若有问题欢迎指正讨论