From 877d07f3213b68f93c01944c1568fa4164318ff6 Mon Sep 17 00:00:00 2001 From: Chuang Wang <2128067684@qq.com> Date: Tue, 24 Mar 2026 12:27:10 +0800 Subject: [PATCH] [feature](point query) Support IN predicate in short-circuit point query --- .../doris/nereids/StatementContext.java | 7 + .../rules/analysis/ExpressionAnalyzer.java | 25 ++ ...calResultSinkToShortCircuitPointQuery.java | 56 ++++- .../apache/doris/qe/PointQueryExecutor.java | 132 ++++++++--- .../rewrite/ShortCircuitInQueryTest.java | 223 ++++++++++++++++++ 5 files changed, 397 insertions(+), 46 deletions(-) create mode 100644 fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitInQueryTest.java diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java index af24f9a83c7dfe..10d246fb296a19 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java @@ -177,6 +177,9 @@ public enum TableFrom { // map placeholder id to comparison slot, which will used to replace conjuncts // directly private final Map idToComparisonSlot = new TreeMap<>(); + // map placeholder id to slot for IN predicate options, used to replace IN predicate + // conjuncts in short circuit plan for prepared statement + private final Map idToInPredicateSlot = new TreeMap<>(); // collect all hash join conditions to compute node connectivity in join graph private final List joinFilters = new ArrayList<>(); @@ -628,6 +631,10 @@ public Map getIdToComparisonSlot() { return idToComparisonSlot; } + public Map getIdToInPredicateSlot() { + return idToInPredicateSlot; + } + public Map, Group>>> getCteIdToConsumerGroup() { return cteIdToConsumerGroup; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java index 3d1faf318e6408..d31a33eed96a97 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java @@ -893,12 +893,37 @@ public Expression visitWhenClause(WhenClause whenClause, ExpressionRewriteContex @Override public Expression visitInPredicate(InPredicate inPredicate, ExpressionRewriteContext context) { + // Register placeholder ids to slot BEFORE children are visited (i.e., before Placeholders + // are replaced with actual literal values by visitPlaceholder). + // Used to replace expressions in ShortCircuit plan for prepared statement. + registerInPredicatePlaceholderToSlot(inPredicate, context); List rewrittenChildren = inPredicate.children().stream() .map(e -> e.accept(this, context)).collect(Collectors.toList()); InPredicate newInPredicate = inPredicate.withChildren(rewrittenChildren); return TypeCoercionUtils.processInPredicate(newInPredicate); } + // Register prepared statement placeholder ids to related slot in IN predicate. + // Each placeholder in the IN list is mapped to the compare slot for short circuit plan replacement. + // Must be called BEFORE children are recursively visited (before Placeholder→Literal substitution). + private void registerInPredicatePlaceholderToSlot(InPredicate inPredicate, + ExpressionRewriteContext context) { + if (context == null) { + return; + } + if (ConnectContext.get() != null + && ConnectContext.get().getCommand() == MysqlCommand.COM_STMT_EXECUTE + && inPredicate.getCompareExpr() instanceof SlotReference) { + SlotReference slot = (SlotReference) inPredicate.getCompareExpr(); + for (Expression option : inPredicate.getOptions()) { + if (option instanceof Placeholder) { + PlaceholderId id = ((Placeholder) option).getPlaceholderId(); + context.cascadesContext.getStatementContext().getIdToInPredicateSlot().put(id, slot); + } + } + } + } + @Override public Expression visitBetween(Between between, ExpressionRewriteContext context) { Expression compareExpr = between.getCompareExpr().accept(this, context); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java index dfcd2ea289cc6f..108c542e682af4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java @@ -25,6 +25,9 @@ import org.apache.doris.nereids.trees.expressions.Cast; import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.InPredicate; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.expressions.Placeholder; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; @@ -39,7 +42,7 @@ /** * short circuit query optimization - * pattern : select xxx from tbl where key = ? + * pattern : select xxx from tbl where key = ? or key IN (?, ?, ...) */ public class LogicalResultSinkToShortCircuitPointQuery implements RewriteRuleFactory { @@ -50,14 +53,35 @@ private Expression removeCast(Expression expression) { return expression; } - private boolean filterMatchShortCircuitCondition(LogicalFilter filter) { - return filter.getConjuncts().stream().allMatch( - // all conjuncts match with pattern `key = ?` - expression -> (expression instanceof EqualTo) - && (removeCast(expression.child(0)).isKeyColumnFromTable() + // Check if an expression in the filter is a valid short-circuit condition. + // Supports: key = literal/placeholder, or key IN (literal/placeholder, ...) + private boolean isValidShortCircuitExpression(Expression expression) { + if (expression instanceof EqualTo) { + // key = literal or key = placeholder + return (removeCast(expression.child(0)).isKeyColumnFromTable() || (expression.child(0) instanceof SlotReference && ((SlotReference) expression.child(0)).getName().equals(Column.DELETE_SIGN))) - && expression.child(1).isLiteral()); + && (expression.child(1).isLiteral() || expression.child(1) instanceof Placeholder); + } else if (expression instanceof InPredicate) { + // key IN (literal/placeholder, ...) + InPredicate inPredicate = (InPredicate) expression; + Expression compareExpr = removeCast(inPredicate.getCompareExpr()); + if (!compareExpr.isKeyColumnFromTable()) { + return false; + } + // All options must be literals or placeholders + for (Expression option : inPredicate.getOptions()) { + if (!(option instanceof Literal) && !(option instanceof Placeholder)) { + return false; + } + } + return !inPredicate.getOptions().isEmpty(); + } + return false; + } + + private boolean filterMatchShortCircuitCondition(LogicalFilter filter) { + return filter.getConjuncts().stream().allMatch(this::isValidShortCircuitExpression); } private boolean scanMatchShortCircuitCondition(LogicalOlapScan olapScan) { @@ -75,12 +99,21 @@ private boolean scanMatchShortCircuitCondition(LogicalOlapScan olapScan) { // set short circuit flag and return the original plan private Plan shortCircuit(Plan root, OlapTable olapTable, Set conjuncts, StatementContext statementContext) { - // All key columns in conjuncts + // All key columns covered by conjuncts (EqualTo or InPredicate) Set colNames = Sets.newHashSet(); for (Expression expr : conjuncts) { - SlotReference slot = ((SlotReference) removeCast((expr.child(0)))); - if (slot.isKeyColumnFromTable()) { - colNames.add(slot.getName()); + if (expr instanceof EqualTo) { + SlotReference slot = (SlotReference) removeCast(expr.child(0)); + if (slot.isKeyColumnFromTable()) { + colNames.add(slot.getName()); + } + } else if (expr instanceof InPredicate) { + InPredicate inPredicate = (InPredicate) expr; + Expression compareExpr = removeCast(inPredicate.getCompareExpr()); + if (compareExpr instanceof SlotReference + && ((SlotReference) compareExpr).isKeyColumnFromTable()) { + colNames.add(((SlotReference) compareExpr).getName()); + } } } // set short circuit flag and modify nothing to the plan @@ -99,7 +132,6 @@ public List buildRules() { ).when(this::filterMatchShortCircuitCondition))) .thenApply(ctx -> { return shortCircuit(ctx.root, ctx.root.child().child().child().getTable(), - ctx.root.child().child().getConjuncts(), ctx.statementContext); })), RuleType.SHOR_CIRCUIT_POINT_QUERY.build( diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/PointQueryExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/PointQueryExecutor.java index 01559bbb64dbd1..f2c31bae2d7b09 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/PointQueryExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/PointQueryExecutor.java @@ -20,6 +20,7 @@ import org.apache.doris.analysis.BinaryPredicate; import org.apache.doris.analysis.Expr; import org.apache.doris.analysis.ExprToSqlVisitor; +import org.apache.doris.analysis.InPredicate; import org.apache.doris.analysis.LiteralExpr; import org.apache.doris.analysis.SlotRef; import org.apache.doris.analysis.ToSqlParams; @@ -145,7 +146,7 @@ public static void directExecuteShortCircuitQuery(StmtExecutor executor, StatementContext statementContext) throws Exception { Preconditions.checkNotNull(preparedStmtCtx.shortCircuitQueryContext); ShortCircuitQueryContext shortCircuitQueryContext = preparedStmtCtx.shortCircuitQueryContext.get(); - // update conjuncts + // update conjuncts for equality predicates: colName -> LiteralExpr Map colNameToConjunct = Maps.newHashMap(); for (Entry entry : statementContext.getIdToComparisonSlot().entrySet()) { String colName = entry.getValue().getOriginalColumn().get().getName(); @@ -153,13 +154,24 @@ public static void directExecuteShortCircuitQuery(StmtExecutor executor, .get(entry.getKey())).toLegacyLiteral(); colNameToConjunct.put(colName, conjunctVal); } - if (colNameToConjunct.size() != preparedStmtCtx.command.placeholderCount()) { + // update conjuncts for IN predicate: colName -> List + Map> colNameToInValues = Maps.newHashMap(); + for (Entry entry : statementContext.getIdToInPredicateSlot().entrySet()) { + String colName = entry.getValue().getOriginalColumn().get().getName(); + Expr inVal = ((Literal) statementContext.getIdToPlaceholderRealExpr() + .get(entry.getKey())).toLegacyLiteral(); + colNameToInValues.computeIfAbsent(colName, k -> new ArrayList<>()).add(inVal); + } + int totalPlaceholderCount = preparedStmtCtx.command.placeholderCount(); + int resolvedCount = colNameToConjunct.size() + colNameToInValues.values().stream() + .mapToInt(List::size).sum(); + if (resolvedCount != totalPlaceholderCount) { throw new AnalysisException("Mismatched conjuncts values size with prepared" + "statement parameters size, expected " - + preparedStmtCtx.command.placeholderCount() - + ", but meet " + colNameToConjunct.size()); + + totalPlaceholderCount + + ", but meet " + resolvedCount); } - updateScanNodeConjuncts(shortCircuitQueryContext.scanNode, colNameToConjunct); + updateScanNodeConjuncts(shortCircuitQueryContext.scanNode, colNameToConjunct, colNameToInValues); // short circuit plan and execution executor.executeAndSendResult(false, false, shortCircuitQueryContext.analzyedQuery, executor.getContext() @@ -167,25 +179,41 @@ public static void directExecuteShortCircuitQuery(StmtExecutor executor, } private static void updateScanNodeConjuncts(OlapScanNode scanNode, - Map colNameToConjunct) { + Map colNameToConjunct, Map> colNameToInValues) { for (Expr conjunct : scanNode.getConjuncts()) { - BinaryPredicate binaryPredicate = (BinaryPredicate) conjunct; - SlotRef slot = null; - int updateChildIdx = 0; - if (binaryPredicate.getChild(0) instanceof LiteralExpr) { - slot = (SlotRef) binaryPredicate.getChildWithoutCast(1); - } else if (binaryPredicate.getChild(1) instanceof LiteralExpr) { - slot = (SlotRef) binaryPredicate.getChildWithoutCast(0); - updateChildIdx = 1; - } else { - Preconditions.checkState(false, "Should contains literal in " - + binaryPredicate.accept(ExprToSqlVisitor.INSTANCE, ToSqlParams.WITH_TABLE)); - } - // not a placeholder to replace - if (!colNameToConjunct.containsKey(slot.getColumnName())) { - continue; + if (conjunct instanceof BinaryPredicate) { + BinaryPredicate binaryPredicate = (BinaryPredicate) conjunct; + SlotRef slot = null; + int updateChildIdx = 0; + if (binaryPredicate.getChild(0) instanceof LiteralExpr) { + slot = (SlotRef) binaryPredicate.getChildWithoutCast(1); + } else if (binaryPredicate.getChild(1) instanceof LiteralExpr) { + slot = (SlotRef) binaryPredicate.getChildWithoutCast(0); + updateChildIdx = 1; + } else { + Preconditions.checkState(false, "Should contains literal in " + + binaryPredicate.accept(ExprToSqlVisitor.INSTANCE, ToSqlParams.WITH_TABLE)); + } + // not a placeholder to replace + if (!colNameToConjunct.containsKey(slot.getColumnName())) { + continue; + } + binaryPredicate.setChild(updateChildIdx, colNameToConjunct.get(slot.getColumnName())); + } else if (conjunct instanceof InPredicate) { + InPredicate inPredicate = (InPredicate) conjunct; + SlotRef slot = inPredicate.getChild(0).unwrapSlotRef(); + if (slot == null || !colNameToInValues.containsKey(slot.getColumnName())) { + continue; + } + List newValues = colNameToInValues.get(slot.getColumnName()); + // Replace all list children (children[1..n]) with the new literal values + while (inPredicate.getChildren().size() > 1) { + inPredicate.getChildren().remove(inPredicate.getChildren().size() - 1); + } + for (Expr val : newValues) { + inPredicate.addChild(val); + } } - binaryPredicate.setChild(updateChildIdx, colNameToConjunct.get(slot.getColumnName())); } } @@ -195,21 +223,57 @@ public void setTimeout(long timeoutMs) { void addKeyTuples( InternalService.PTabletKeyLookupRequest.Builder requestBuilder) { - // TODO handle IN predicates - Map columnExpr = Maps.newHashMap(); - KeyTuple.Builder kBuilder = KeyTuple.newBuilder(); + // Separate equality predicates from IN predicates + Map equalityColumnExpr = Maps.newHashMap(); + Map> inColumnExprs = Maps.newHashMap(); + for (Expr expr : shortCircuitQueryContext.scanNode.getConjuncts()) { - BinaryPredicate predicate = (BinaryPredicate) expr; - Expr left = predicate.getChild(0); - Expr right = predicate.getChild(1); - SlotRef columnSlot = left.unwrapSlotRef(); - columnExpr.put(columnSlot.getColumnName(), right); + if (expr instanceof BinaryPredicate) { + BinaryPredicate predicate = (BinaryPredicate) expr; + Expr left = predicate.getChild(0); + Expr right = predicate.getChild(1); + SlotRef columnSlot = left.unwrapSlotRef(); + equalityColumnExpr.put(columnSlot.getColumnName(), right); + } else if (expr instanceof InPredicate) { + InPredicate inPredicate = (InPredicate) expr; + SlotRef slot = inPredicate.getChild(0).unwrapSlotRef(); + if (slot != null) { + List values = new ArrayList<>(); + for (int i = 1; i < inPredicate.getChildren().size(); i++) { + values.add(inPredicate.getChild(i)); + } + inColumnExprs.put(slot.getColumnName(), values); + } + } } - // add key tuple in keys order - for (Column column : shortCircuitQueryContext.scanNode.getOlapTable().getBaseSchemaKeyColumns()) { - kBuilder.addKeyColumnRep(columnExpr.get(column.getName()).getStringValue()); + + List keyColumns = shortCircuitQueryContext.scanNode.getOlapTable().getBaseSchemaKeyColumns(); + + if (inColumnExprs.isEmpty()) { + // Pure equality case: generate one KeyTuple + KeyTuple.Builder kBuilder = KeyTuple.newBuilder(); + for (Column column : keyColumns) { + kBuilder.addKeyColumnRep(equalityColumnExpr.get(column.getName()).getStringValue()); + } + requestBuilder.addKeyTuples(kBuilder); + } else { + // IN predicate case: generate one KeyTuple per combination + // Find the IN column and its values, combine with equality columns + // Note: currently only supports single IN column per query (all keys must be covered) + String inColName = inColumnExprs.keySet().iterator().next(); + List inValues = inColumnExprs.get(inColName); + for (Expr inVal : inValues) { + KeyTuple.Builder kBuilder = KeyTuple.newBuilder(); + for (Column column : keyColumns) { + if (column.getName().equals(inColName)) { + kBuilder.addKeyColumnRep(inVal.getStringValue()); + } else { + kBuilder.addKeyColumnRep(equalityColumnExpr.get(column.getName()).getStringValue()); + } + } + requestBuilder.addKeyTuples(kBuilder); + } } - requestBuilder.addKeyTuples(kBuilder); } @Override diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitInQueryTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitInQueryTest.java new file mode 100644 index 00000000000000..98f687d6ee8e17 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/ShortCircuitInQueryTest.java @@ -0,0 +1,223 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.rules.rewrite; + +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.util.PlanChecker; +import org.apache.doris.utframe.TestWithFeService; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Tests that key-column IN queries on MOW tables trigger the short-circuit point query path. + */ +public class ShortCircuitInQueryTest extends TestWithFeService { + + private static final String MOW_TABLE = "mow_tbl"; + private static final String SINGLE_KEY_TABLE = "mow_single_key"; + private static final String DUP_TABLE = "dup_tbl"; + + @Override + protected void runBeforeAll() throws Exception { + createDatabase("test"); + connectContext.setDatabase("test"); + connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + + // MOW table with two key columns, light_schema_change and store_row_column enabled. + // All three properties are required by scanMatchShortCircuitCondition. + createTables( + "CREATE TABLE IF NOT EXISTS " + MOW_TABLE + " (\n" + + " k1 INT NOT NULL,\n" + + " k2 INT NOT NULL,\n" + + " v1 VARCHAR(100)\n" + + ") ENGINE=OLAP\n" + + "UNIQUE KEY(k1, k2)\n" + + "DISTRIBUTED BY HASH(k1) BUCKETS 4\n" + + "PROPERTIES (\n" + + " \"replication_num\" = \"1\",\n" + + " \"enable_unique_key_merge_on_write\" = \"true\",\n" + + " \"light_schema_change\" = \"true\",\n" + + " \"store_row_column\" = \"true\"\n" + + ")", + + // MOW table with a single key column, used to verify IN on the only key. + "CREATE TABLE IF NOT EXISTS " + SINGLE_KEY_TABLE + " (\n" + + " id INT NOT NULL,\n" + + " name VARCHAR(100)\n" + + ") ENGINE=OLAP\n" + + "UNIQUE KEY(id)\n" + + "DISTRIBUTED BY HASH(id) BUCKETS 4\n" + + "PROPERTIES (\n" + + " \"replication_num\" = \"1\",\n" + + " \"enable_unique_key_merge_on_write\" = \"true\",\n" + + " \"light_schema_change\" = \"true\",\n" + + " \"store_row_column\" = \"true\"\n" + + ")", + + // Duplicate-key table – must NOT trigger short-circuit even with IN on key columns. + "CREATE TABLE IF NOT EXISTS " + DUP_TABLE + " (\n" + + " k1 INT NOT NULL,\n" + + " k2 INT NOT NULL,\n" + + " v1 VARCHAR(100)\n" + + ") ENGINE=OLAP\n" + + "DUPLICATE KEY(k1, k2)\n" + + "DISTRIBUTED BY HASH(k1) BUCKETS 4\n" + + "PROPERTIES (\n" + + " \"replication_num\" = \"1\"\n" + + ")" + ); + } + + // ----------------------------------------------------------------------- + // Positive cases: should trigger short-circuit (isShortCircuitQuery = true) + // ----------------------------------------------------------------------- + + /** + * SELECT * FROM mow_single_key WHERE id IN (1, 2, 3) + * Only one key column, covered by IN → short-circuit. + */ + @Test + void testSingleKeyInQueryTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + SINGLE_KEY_TABLE + " WHERE id IN (1, 2, 3)") + .rewrite() + .getCascadesContext(); + Assertions.assertTrue(ctx.getStatementContext().isShortCircuitQuery(), + "Single-key IN query on MOW table should trigger short-circuit"); + } + + /** + * SELECT * FROM mow_single_key WHERE id IN (42) + * IN with a single element is still a valid IN predicate → short-circuit. + */ + @Test + void testSingleKeyInWithOneValueTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + SINGLE_KEY_TABLE + " WHERE id IN (42)") + .rewrite() + .getCascadesContext(); + Assertions.assertTrue(ctx.getStatementContext().isShortCircuitQuery(), + "IN with a single value on the only key column should trigger short-circuit"); + } + + /** + * SELECT * FROM mow_tbl WHERE k1 IN (1, 2) AND k2 = 10 + * Two key columns: k1 covered by IN, k2 by equality → all keys covered → short-circuit. + */ + @Test + void testCompositeKeyInAndEqualTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + MOW_TABLE + + " WHERE k1 IN (1, 2) AND k2 = 10") + .rewrite() + .getCascadesContext(); + Assertions.assertTrue(ctx.getStatementContext().isShortCircuitQuery(), + "IN on first key + equality on second key should trigger short-circuit"); + } + + /** + * SELECT * FROM mow_tbl WHERE k1 = 5 AND k2 IN (10, 20, 30) + * Same as above with the roles swapped. + */ + @Test + void testCompositeKeyEqualAndInTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + MOW_TABLE + + " WHERE k1 = 5 AND k2 IN (10, 20, 30)") + .rewrite() + .getCascadesContext(); + Assertions.assertTrue(ctx.getStatementContext().isShortCircuitQuery(), + "Equality on first key + IN on second key should trigger short-circuit"); + } + + /** + * SELECT * FROM mow_tbl WHERE k1 = 1 AND k2 = 2 + * Pure equality (original behaviour) must still work. + */ + @Test + void testPureEqualityStillTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + MOW_TABLE + " WHERE k1 = 1 AND k2 = 2") + .rewrite() + .getCascadesContext(); + Assertions.assertTrue(ctx.getStatementContext().isShortCircuitQuery(), + "Pure equality on all key columns should still trigger short-circuit"); + } + + // ----------------------------------------------------------------------- + // Negative cases: should NOT trigger short-circuit + // ----------------------------------------------------------------------- + + /** + * SELECT * FROM dup_tbl WHERE k1 IN (1, 2, 3) AND k2 = 1 + * Duplicate-key table: store_row_column / light_schema_change not set → no short-circuit. + */ + @Test + void testDuplicateKeyTableNotTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + DUP_TABLE + " WHERE k1 IN (1, 2, 3) AND k2 = 1") + .rewrite() + .getCascadesContext(); + Assertions.assertFalse(ctx.getStatementContext().isShortCircuitQuery(), + "Duplicate-key table should NOT trigger short-circuit"); + } + + /** + * SELECT * FROM mow_tbl WHERE k1 IN (1, 2, 3) + * Only k1 is provided via IN; k2 is missing → key columns not fully covered → no short-circuit. + */ + @Test + void testPartialKeyInNotTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + MOW_TABLE + " WHERE k1 IN (1, 2, 3)") + .rewrite() + .getCascadesContext(); + Assertions.assertFalse(ctx.getStatementContext().isShortCircuitQuery(), + "IN on only part of the key columns should NOT trigger short-circuit"); + } + + /** + * SELECT * FROM mow_tbl WHERE v1 IN ('a', 'b') AND k1 = 1 AND k2 = 2 + * IN on a non-key (value) column v1 is NOT a valid short-circuit expression → no short-circuit. + */ + @Test + void testNonKeyColumnInNotTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + MOW_TABLE + + " WHERE v1 IN ('a', 'b') AND k1 = 1 AND k2 = 2") + .rewrite() + .getCascadesContext(); + Assertions.assertFalse(ctx.getStatementContext().isShortCircuitQuery(), + "IN on a non-key column should NOT trigger short-circuit"); + } + + /** + * SELECT * FROM mow_tbl WHERE k1 > 5 AND k2 = 2 + * Range predicate (>) on a key column is not a valid short-circuit expression → no short-circuit. + */ + @Test + void testRangePredicateNotTriggerShortCircuit() { + CascadesContext ctx = PlanChecker.from(connectContext) + .analyze("SELECT * FROM " + MOW_TABLE + " WHERE k1 > 5 AND k2 = 2") + .rewrite() + .getCascadesContext(); + Assertions.assertFalse(ctx.getStatementContext().isShortCircuitQuery(), + "Range predicate on key column should NOT trigger short-circuit"); + } +}