Examples with ASTNode com.alibaba.cobar.parser.ast.ASTNode used on opensource projects

Example 1 with ASTNode

use of com.alibaba.cobar.parser.ast.ASTNode in project cobar by alibaba.

the class ServerRouter method replacePartitionKeyOperand.

private static void replacePartitionKeyOperand(Map<String, Map<Object, Set<Pair<Expression, ASTNode>>>> index, List<String> cols) {
    if (cols == null) {
        return;
    }
    for (String col : cols) {
        Map<Object, Set<Pair<Expression, ASTNode>>> map = index.get(col);
        if (map == null) {
            continue;
        }
        for (Set<Pair<Expression, ASTNode>> set : map.values()) {
            if (set == null) {
                continue;
            }
            for (Pair<Expression, ASTNode> p : set) {
                Expression expr = p.getKey();
                ASTNode parent = p.getValue();
                if (PartitionKeyVisitor.isPartitionKeyOperandSingle(expr, parent)) {
                    ((ReplacableExpression) expr).setReplaceExpr(ReplacableExpression.BOOL_FALSE);
                } else if (PartitionKeyVisitor.isPartitionKeyOperandIn(expr, parent)) {
                    ((ReplacableExpression) parent).setReplaceExpr(ReplacableExpression.BOOL_FALSE);
                }
            }
        }
    }
}

Example 2 with ASTNode

use of com.alibaba.cobar.parser.ast.ASTNode in project cobar by alibaba.

the class ServerRouter method dispatchInsertReplace.

@SuppressWarnings("unchecked")
private static void dispatchInsertReplace(RouteResultsetNode[] rn, DMLInsertReplaceStatement stmt, List<String> ruleColumns, Map<Integer, List<Object[]>> dataNodeMap, TableConfig matchedTable, String originalSQL, PartitionKeyVisitor visitor) {
    if (stmt.getSelect() != null) {
        dispatchWhereBasedStmt(rn, stmt, ruleColumns, dataNodeMap, matchedTable, originalSQL, visitor);
        return;
    }
    Map<String, Map<Object, Set<Pair<Expression, ASTNode>>>> colsIndex = visitor.getColumnIndex(stmt.getTable().getIdTextUpUnescape());
    if (colsIndex == null || colsIndex.isEmpty()) {
        throw new IllegalArgumentException("columns index is empty: " + originalSQL);
    }
    ArrayList<Map<Object, Set<Pair<Expression, ASTNode>>>> colsIndexList = new ArrayList<Map<Object, Set<Pair<Expression, ASTNode>>>>(ruleColumns.size());
    for (int i = 0, len = ruleColumns.size(); i < len; ++i) {
        colsIndexList.add(colsIndex.get(ruleColumns.get(i)));
    }
    int dataNodeId = -1;
    for (Entry<Integer, List<Object[]>> en : dataNodeMap.entrySet()) {
        List<Object[]> tuples = en.getValue();
        HashSet<RowExpression> replaceRowList = new HashSet<RowExpression>(tuples.size());
        for (Object[] tuple : tuples) {
            Set<Pair<Expression, ASTNode>> tupleExprs = null;
            for (int i = 0; i < tuple.length; ++i) {
                Map<Object, Set<Pair<Expression, ASTNode>>> valueMap = colsIndexList.get(i);
                Object value = tuple[i];
                Set<Pair<Expression, ASTNode>> set = getExpressionSet(valueMap, value);
                tupleExprs = (Set<Pair<Expression, ASTNode>>) CollectionUtil.intersectSet(tupleExprs, set);
            }
            if (tupleExprs == null || tupleExprs.isEmpty()) {
                throw new IllegalArgumentException("route: empty expression list for insertReplace stmt: " + originalSQL);
            }
            for (Pair<Expression, ASTNode> p : tupleExprs) {
                if (p.getValue() == stmt && p.getKey() instanceof RowExpression) {
                    replaceRowList.add((RowExpression) p.getKey());
                }
            }
        }
        stmt.setReplaceRowList(new ArrayList<RowExpression>(replaceRowList));
        String sql = genSQL(stmt, originalSQL);
        stmt.clearReplaceRowList();
        String dataNodeName = matchedTable.getDataNodes()[en.getKey()];
        rn[++dataNodeId] = new RouteResultsetNode(dataNodeName, sql);
    }
}

Example 3 with ASTNode

use of com.alibaba.cobar.parser.ast.ASTNode in project cobar by alibaba.

the class ServerRouter method dispatchWhereBasedStmt.

private static void dispatchWhereBasedStmt(RouteResultsetNode[] rn, SQLStatement stmtAST, List<String> ruleColumns, Map<Integer, List<Object[]>> dataNodeMap, TableConfig matchedTable, String originalSQL, PartitionKeyVisitor visitor) {
    // [perf tag] 11.617 us: sharding multivalue
    if (ruleColumns.size() > 1) {
        String sql;
        if (visitor.isSchemaTrimmed()) {
            sql = genSQL(stmtAST, originalSQL);
        } else {
            sql = originalSQL;
        }
        int i = -1;
        for (Integer dataNodeId : dataNodeMap.keySet()) {
            String dataNode = matchedTable.getDataNodes()[dataNodeId];
            rn[++i] = new RouteResultsetNode(dataNode, sql);
        }
        return;
    }
    final String table = matchedTable.getName();
    Map<String, Map<Object, Set<Pair<Expression, ASTNode>>>> columnIndex = visitor.getColumnIndex(table);
    Map<Object, Set<Pair<Expression, ASTNode>>> valueMap = columnIndex.get(ruleColumns.get(0));
    replacePartitionKeyOperand(columnIndex, ruleColumns);
    Map<InExpression, Set<Expression>> unreplacedInExpr = new HashMap<InExpression, Set<Expression>>(1, 1);
    Set<ReplacableExpression> unreplacedSingleExprs = new HashSet<ReplacableExpression>();
    // [perf tag] 12.2755 us: sharding multivalue
    int nodeId = -1;
    for (Entry<Integer, List<Object[]>> en : dataNodeMap.entrySet()) {
        List<Object[]> tuples = en.getValue();
        unreplacedSingleExprs.clear();
        unreplacedInExpr.clear();
        for (Object[] tuple : tuples) {
            Object value = tuple[0];
            Set<Pair<Expression, ASTNode>> indexedExpressionPair = getExpressionSet(valueMap, value);
            for (Pair<Expression, ASTNode> pair : indexedExpressionPair) {
                Expression expr = pair.getKey();
                ASTNode parent = pair.getValue();
                if (PartitionKeyVisitor.isPartitionKeyOperandSingle(expr, parent)) {
                    unreplacedSingleExprs.add((ReplacableExpression) expr);
                } else if (PartitionKeyVisitor.isPartitionKeyOperandIn(expr, parent)) {
                    Set<Expression> newInSet = unreplacedInExpr.get(parent);
                    if (newInSet == null) {
                        newInSet = new HashSet<Expression>(indexedExpressionPair.size(), 1);
                        unreplacedInExpr.put((InExpression) parent, newInSet);
                    }
                    newInSet.add(expr);
                }
            }
        }
        for (ReplacableExpression expr : unreplacedSingleExprs) {
            expr.clearReplaceExpr();
        }
        for (Entry<InExpression, Set<Expression>> entemp : unreplacedInExpr.entrySet()) {
            InExpression in = entemp.getKey();
            Set<Expression> set = entemp.getValue();
            if (set == null || set.isEmpty()) {
                in.setReplaceExpr(ReplacableExpression.BOOL_FALSE);
            } else {
                in.clearReplaceExpr();
                InExpressionList inlist = in.getInExpressionList();
                if (inlist != null)
                    inlist.setReplaceExpr(new ArrayList<Expression>(set));
            }
        }
        // [perf tag] 16.506 us: sharding multivalue
        String sql = genSQL(stmtAST, originalSQL);
        // [perf tag] 21.3425 us: sharding multivalue
        String dataNodeName = matchedTable.getDataNodes()[en.getKey()];
        rn[++nodeId] = new RouteResultsetNode(dataNodeName, sql);
        for (ReplacableExpression expr : unreplacedSingleExprs) {
            expr.setReplaceExpr(ReplacableExpression.BOOL_FALSE);
        }
        for (InExpression in : unreplacedInExpr.keySet()) {
            in.setReplaceExpr(ReplacableExpression.BOOL_FALSE);
            InExpressionList list = in.getInExpressionList();
            if (list != null)
                list.clearReplaceExpr();
        }
    // [perf tag] 22.0965 us: sharding multivalue
    }
}

Example 4 with ASTNode

use of com.alibaba.cobar.parser.ast.ASTNode in project cobar by alibaba.

the class PartitionKeyVisitor method insertReplace.

private void insertReplace(DMLInsertReplaceStatement node) {
    Identifier table = node.getTable();
    List<Identifier> collist = node.getColumnNameList();
    QueryExpression query = node.getSelect();
    List<RowExpression> rows = node.getRowList();
    tableAsTableFactor(table);
    String tableName = table.getIdTextUpUnescape();
    visitChild(2, false, false, collist);
    if (query != null) {
        query.accept(this);
        return;
    }
    for (RowExpression row : rows) {
        visitChild(2, false, false, row);
    }
    Map<String, List<Object>> colVals = ensureColumnValueByTable(tableName);
    Map<String, Map<Object, Set<Pair<Expression, ASTNode>>>> colValsIndex = ensureColumnValueIndexByTable(tableName);
    if (collist != null) {
        for (int i = 0; i < collist.size(); ++i) {
            String colName = collist.get(i).getIdTextUpUnescape();
            if (isRuledColumn(tableName, colName)) {
                List<Object> valueList = ensureColumnValueList(colVals, colName);
                Map<Object, Set<Pair<Expression, ASTNode>>> valMap = ensureColumnValueIndexObjMap(colValsIndex, colName);
                for (RowExpression row : rows) {
                    Expression expr = row.getRowExprList().get(i);
                    Object value = expr == null ? null : expr.evaluation(evaluationParameter);
                    if (value != Expression.UNEVALUATABLE) {
                        valueList.add(value);
                        addIntoColumnValueIndex(valMap, value, row, node);
                    }
                }
            }
        }
    }
}

Example 5 with ASTNode

use of com.alibaba.cobar.parser.ast.ASTNode in project cobar by alibaba.

the class PartitionKeyVisitor method visit.

@Override
public void visit(InExpression node) {
    Expression left = node.getLeftOprand();
    Expression right = node.getRightOprand();
    visitChild(2, false, false, left, right);
    if (verdictColumn && !node.isNot() && left instanceof Identifier && right instanceof InExpressionList) {
        Identifier col = (Identifier) left;
        String colName = col.getIdTextUpUnescape();
        String table = tableAlias.get(col.getLevelUnescapeUpName(2));
        if (isRuledColumn(table, colName)) {
            List<Object> valList = ensureColumnValueList(ensureColumnValueByTable(table), colName);
            Map<Object, Set<Pair<Expression, ASTNode>>> valMap = ensureColumnValueIndexObjMap(ensureColumnValueIndexByTable(table), colName);
            InExpressionList inlist = (InExpressionList) right;
            for (Expression expr : inlist.getList()) {
                Object value = expr.evaluation(evaluationParameter);
                if (value != Expression.UNEVALUATABLE) {
                    valList.add(value);
                    addIntoColumnValueIndex(valMap, value, expr, node);
                }
            }
        }
    }
}