eldenmoon commented on code in PR #39468: URL: https://github.com/apache/doris/pull/39468#discussion_r1746499775
########## fe/fe-core/src/main/java/org/apache/doris/qe/PointQueryExecutor.java: ########## @@ -159,23 +191,181 @@ public static void directExecuteShortCircuitQuery(StmtExecutor executor, .getMysqlChannel(), null, null); } - private static void updateScanNodeConjuncts(OlapScanNode scanNode, List<Expr> conjunctVals) { - for (int i = 0; i < conjunctVals.size(); ++i) { - BinaryPredicate binaryPredicate = (BinaryPredicate) scanNode.getConjuncts().get(i); - if (binaryPredicate.getChild(0) instanceof LiteralExpr) { - binaryPredicate.setChild(0, conjunctVals.get(i)); - } else if (binaryPredicate.getChild(1) instanceof LiteralExpr) { - binaryPredicate.setChild(1, conjunctVals.get(i)); + /* + * Interface for handling different conjunct types + */ + private interface ConjunctHandler { + int handle(Expr expr, List<Expr> conjunctVals, int conjunctValsBegin) throws AnalysisException; + } + + private static class InPredicateHandler implements ConjunctHandler { + @Override + public int handle(Expr expr, List<Expr> conjunctVals, int conjunctValsBegin) throws AnalysisException { + InPredicate inPredicate = (InPredicate) expr; + if (inPredicate.isNotIn()) { + throw new AnalysisException("Not support NOT IN predicate in point query"); + } + for (int j = 1; j < inPredicate.getChildren().size(); ++j) { + if (inPredicate.getChild(j) instanceof LiteralExpr) { + inPredicate.setChild(j, conjunctVals.get(conjunctValsBegin++)); + } else { + Preconditions.checkState(false, "Should contains literal in " + inPredicate.toSqlImpl()); + } + } + return conjunctValsBegin; + } + } + + private static class BinaryPredicateHandler implements ConjunctHandler { + @Override + public int handle(Expr expr, List<Expr> conjunctVals, int conjunctValsBegin) throws AnalysisException { + BinaryPredicate binaryPredicate = (BinaryPredicate) expr; + Expr left = binaryPredicate.getChild(0); + Expr right = binaryPredicate.getChild(1); + + if (isDeleteSign(left) || isDeleteSign(right)) { + return conjunctValsBegin; + } + + if (isLiteralExpr(left)) { + binaryPredicate.setChild(0, conjunctVals.get(conjunctValsBegin++)); + } else if (isLiteralExpr(right)) { + binaryPredicate.setChild(1, conjunctVals.get(conjunctValsBegin++)); } else { Preconditions.checkState(false, "Should contains literal in " + binaryPredicate.toSqlImpl()); } + return conjunctValsBegin; } + + private boolean isLiteralExpr(Expr expr) { + return expr instanceof LiteralExpr; + } + + private boolean isDeleteSign(Expr expr) { + return expr instanceof SlotRef && ((SlotRef) expr).getColumnName().equalsIgnoreCase(Column.DELETE_SIGN); + } + } + + private static void updateScanNodeConjuncts(OlapScanNode scanNode, List<Expr> conjunctVals) { + List<Expr> conjuncts = scanNode.getConjuncts(); + Map<Class<? extends Expr>, ConjunctHandler> handlers = Maps.newHashMap(); + handlers.put(InPredicate.class, new InPredicateHandler()); + handlers.put(BinaryPredicate.class, new BinaryPredicateHandler()); + int conjunctValsBegin = 0; + for (Expr expr : conjuncts) { + ConjunctHandler handler = handlers.get(expr.getClass()); + if (handler == null) { + throw new AnalysisException("Not support conjunct type " + expr.getClass().getName()); + } + conjunctValsBegin = handler.handle(expr, conjunctVals, conjunctValsBegin); + } + + Preconditions.checkState(conjunctValsBegin == conjunctVals.size()); + } + + public static byte[] appendBytes(byte[] original, byte[] toAppend) { + byte[] result = new byte[original.length + toAppend.length]; + System.arraycopy(original, 0, result, 0, original.length); + System.arraycopy(toAppend, 0, result, original.length, toAppend.length); + return result; } public void setTimeout(long timeoutMs) { this.timeoutMs = timeoutMs; } + void getAllKeyTupleCombination(List<Expr> conjuncts, int index, + List<String> currentKeyTuple, // current key tuple + List<List<String>> result, + List<String> columnExpr, + List<Column> keyColumns) { + if (index == conjuncts.size()) { + List<String> orderedKeyTuple = new ArrayList<>(currentKeyTuple.size()); + OlapTable olapTable = shortCircuitQueryContext.scanNode.getOlapTable(); + + // add key tuple in keys order + for (Column column : keyColumns) { + int colIdx = columnExpr.indexOf(column.getName()); + String currentKey = currentKeyTuple.get(colIdx); + orderedKeyTuple.add(currentKey); + } + List<String> keyTupleForDistributionPrune = + Lists.newArrayList(distributionKeyColumns.stream().map((idx) -> { + return orderedKeyTuple.get(idx); + }).collect(Collectors.toList())); + result.add(Lists.newArrayList(orderedKeyTuple)); + Set<Long> tabletIDs = Sets.newHashSet(); + for (PartitionKey key : distributionKeys2TabletID.keySet()) { + List<String> distributionKeys = Lists.newArrayList(); + for (LiteralExpr expr : key.getKeys()) { + distributionKeys.add(expr.getStringValue()); + } + if (distributionKeys.equals(keyTupleForDistributionPrune)) { + tabletIDs.addAll(distributionKeys2TabletID.get(key)); + break; + } + } + // if all key columns are distribution columns, then we don't need to prune + if (distributionKeyColumns.size() == keyColumns.size()) { + keyTupleID2TabletID.add(tabletIDs); + return; + } + for (Integer colIdx : partitionKeyColumns) { + String partitionColName = keyColumns.get(colIdx).getName(); + String keyColValue = orderedKeyTuple.get(colIdx); + try { + ColumnBound bound = + ColumnBound.of(LiteralExpr.create(keyColValue, keyColumns.get(colIdx).getType())); + List<Long> partitionID = partitionCol2PartitionID.get(partitionColName) + .get(bound); + Preconditions.checkState(partitionID.size() == 1); + Partition partition = olapTable.getPartition(partitionID.get(0)); + MaterializedIndex selectedTable = + partition.getIndex(shortCircuitQueryContext.scanNode.getSelectedIndexId()); + tabletIDs.retainAll(selectedTable.getTablets() Review Comment: ? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For additional commands, e-mail: commits-h...@doris.apache.org