Skip to content

Commit 0410274

Browse files
trueeyuCopilotclaude
authored
[UT] Add regression test for lambda argument id collision across ColumnRefFactory (StarRocks#75297)
Signed-off-by: trueeyu <lxhhust350@qq.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent d2369a9 commit 0410274

1 file changed

Lines changed: 177 additions & 0 deletions

File tree

‎fe/fe-core/src/test/java/com/starrocks/sql/plan/LambdaArgIdCollisionTest.java‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
// Copyright 2021-present StarRocks, Inc. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package com.starrocks.sql.plan;
16+
17+
import com.starrocks.sql.analyzer.Analyzer;
18+
import com.starrocks.sql.ast.QueryStatement;
19+
import com.starrocks.sql.ast.StatementBase;
20+
import com.starrocks.sql.optimizer.OptExpression;
21+
import com.starrocks.sql.optimizer.base.ColumnRefFactory;
22+
import com.starrocks.sql.optimizer.operator.OperatorType;
23+
import com.starrocks.sql.optimizer.operator.logical.LogicalJoinOperator;
24+
import com.starrocks.sql.optimizer.operator.logical.LogicalProjectOperator;
25+
import com.starrocks.sql.optimizer.operator.scalar.ColumnRefOperator;
26+
import com.starrocks.sql.optimizer.operator.scalar.LambdaFunctionOperator;
27+
import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
28+
import com.starrocks.sql.optimizer.transformer.LogicalPlan;
29+
import com.starrocks.sql.optimizer.transformer.RelationTransformer;
30+
import com.starrocks.sql.parser.SqlParser;
31+
import org.junit.jupiter.api.Assertions;
32+
import org.junit.jupiter.api.BeforeAll;
33+
import org.junit.jupiter.api.Test;
34+
35+
import java.util.ArrayList;
36+
import java.util.HashMap;
37+
import java.util.List;
38+
import java.util.Map;
39+
import java.util.Set;
40+
import java.util.TreeSet;
41+
42+
// Regression test for historical bug fixed by #73273: when the same analyzed AST is
43+
// transformed twice with different ColumnRefFactory instances, lambda-argument ColumnRefOperator
44+
// ids must not leak across factories and collide with scalar column ids.
45+
//
46+
// The original bug cached the lambda argument ColumnRefOperator on the AST node
47+
// (LambdaArgument.transformedOp). A second transform with a fresh factory could then reuse that
48+
// stale id and collide with a newly created scalar column id, corrupting types (e.g. VARCHAR
49+
// transaction_uuid/client_key becoming the array element struct<...>) and producing invalid
50+
// predicates (e.g. struct = varchar join conjunct).
51+
public class LambdaArgIdCollisionTest extends PlanTestNoneDBBase {
52+
@BeforeAll
53+
public static void beforeClass() throws Exception {
54+
PlanTestNoneDBBase.beforeClass();
55+
starRocksAssert.withDatabase("repro_lambda").useDatabase("repro_lambda");
56+
starRocksAssert.withTable("CREATE TABLE `conn_event` (\n" +
57+
" `transaction_uuid` varchar(200) NULL,\n" +
58+
" `client_key` varchar(200) NULL,\n" +
59+
" `request_attributes` ARRAY<STRUCT<name varchar(200), value varchar(200)>> NULL,\n" +
60+
" `errors` ARRAY<STRUCT<code varchar(200), description varchar(200)>> NULL\n" +
61+
") ENGINE=OLAP\n" +
62+
"DUPLICATE KEY(`transaction_uuid`)\n" +
63+
"DISTRIBUTED BY HASH(`transaction_uuid`) BUCKETS 3\n" +
64+
"PROPERTIES (\"replication_num\" = \"1\");");
65+
}
66+
67+
@Test
68+
public void testLambdaArgDoesNotCorruptScalarColumnType() throws Exception {
69+
String sql = "WITH base AS (\n" +
70+
" SELECT transaction_uuid, client_key, errors,\n" +
71+
" array_filter(request_attributes, x -> x.name = 'unitId')[1].value AS room_id,\n" +
72+
" array_filter(request_attributes, x -> x.name = 'ratePlanId')[1].value AS rate_plan_id\n" +
73+
" FROM conn_event\n" +
74+
"),\n" +
75+
"errs AS (\n" +
76+
" SELECT b.transaction_uuid AS transaction_uuid, unnest.code AS code\n" +
77+
" FROM base b CROSS JOIN UNNEST(b.errors) AS unnest\n" +
78+
")\n" +
79+
"SELECT LEFT(c.transaction_uuid, 50) AS tid, LEFT(c.client_key, 50) AS ck,\n" +
80+
" c.room_id, e.code\n" +
81+
"FROM base c INNER JOIN errs e ON e.transaction_uuid = c.transaction_uuid";
82+
83+
StatementBase statement = SqlParser.parse(sql,
84+
connectContext.getSessionVariable().getSqlMode()).get(0);
85+
Analyzer.analyze(statement, connectContext);
86+
QueryStatement query = (QueryStatement) statement;
87+
88+
// First pass: warms the lambda-argument cache on the shared AST (LambdaArgument.transformedOp).
89+
new RelationTransformer(new ColumnRefFactory(), connectContext).transform(query.getQueryRelation());
90+
91+
// Second pass over the SAME AST with a fresh factory, exactly what plan retry / MV
92+
// rewrite do. On buggy code the cached lambda-arg ids leak in and collide here.
93+
ColumnRefFactory factory2 = new ColumnRefFactory();
94+
LogicalPlan plan2 = new RelationTransformer(factory2, connectContext).transform(query.getQueryRelation());
95+
96+
List<ColumnRefOperator> all = new ArrayList<>();
97+
collectColumnRefs(plan2.getRoot(), all);
98+
99+
// BUG signature: one id is owned by both a lambda argument and a normal column.
100+
Map<Integer, ColumnRefOperator> lambdaById = new HashMap<>();
101+
Map<Integer, ColumnRefOperator> scalarById = new HashMap<>();
102+
for (ColumnRefOperator ref : all) {
103+
if (ref.getOpType() == OperatorType.LAMBDA_ARGUMENT) {
104+
lambdaById.put(ref.getId(), ref);
105+
} else {
106+
scalarById.put(ref.getId(), ref);
107+
}
108+
}
109+
Set<Integer> collisions = new TreeSet<>(lambdaById.keySet());
110+
collisions.retainAll(scalarById.keySet());
111+
Assertions.assertTrue(collisions.isEmpty(),
112+
"lambda argument id collides with a scalar column id (lambda-arg id leaked across "
113+
+ "ColumnRefFactory). colliding ids=" + collisions
114+
+ ", lambda=" + lambdaById + ", scalar=" + scalarById);
115+
116+
// And the varchar scalar columns must never be re-typed as struct.
117+
for (ColumnRefOperator ref : all) {
118+
if (("transaction_uuid".equals(ref.getName()) || "client_key".equals(ref.getName()))
119+
&& ref.getOpType() != OperatorType.LAMBDA_ARGUMENT) {
120+
Assertions.assertFalse(ref.getType().isStructType(),
121+
ref.getName() + " should stay VARCHAR but is " + ref.getType()
122+
+ " — lambda arg id collision corrupted its type.");
123+
}
124+
}
125+
}
126+
127+
private static void collectColumnRefs(OptExpression expr, List<ColumnRefOperator> out) {
128+
if (expr.getOp() instanceof LogicalProjectOperator) {
129+
LogicalProjectOperator project = (LogicalProjectOperator) expr.getOp();
130+
for (Map.Entry<ColumnRefOperator, ScalarOperator> e : project.getColumnRefMap().entrySet()) {
131+
collectFromScalar(e.getKey(), out);
132+
collectFromScalar(e.getValue(), out);
133+
}
134+
}
135+
if (expr.getOp() instanceof LogicalJoinOperator) {
136+
collectFromScalar(((LogicalJoinOperator) expr.getOp()).getOnPredicate(), out);
137+
}
138+
collectFromScalar(expr.getOp().getPredicate(), out);
139+
if (expr.getOp().getProjection() != null) {
140+
for (ScalarOperator s : expr.getOp().getProjection().getColumnRefMap().values()) {
141+
collectFromScalar(s, out);
142+
}
143+
for (ColumnRefOperator c : expr.getOp().getProjection().getColumnRefMap().keySet()) {
144+
collectFromScalar(c, out);
145+
}
146+
}
147+
for (OptExpression child : expr.getInputs()) {
148+
collectColumnRefs(child, out);
149+
}
150+
}
151+
152+
private static void collectFromScalar(ScalarOperator s, List<ColumnRefOperator> out) {
153+
if (s == null) {
154+
return;
155+
}
156+
if (s instanceof ColumnRefOperator) {
157+
out.add((ColumnRefOperator) s);
158+
}
159+
if (s instanceof LambdaFunctionOperator) {
160+
LambdaFunctionOperator lambda = (LambdaFunctionOperator) s;
161+
for (ColumnRefOperator ref : lambda.getRefColumns()) {
162+
out.add(ref);
163+
}
164+
// The CSE/reuse columnRefMap holds ColumnRefOperators (keys) and their defining
165+
// expressions (values) that live only in the map: the lambda body references the
166+
// keys, so the values are not reachable via getChildren()/getLambdaExpr().
167+
for (Map.Entry<ColumnRefOperator, ScalarOperator> e : lambda.getColumnRefMap().entrySet()) {
168+
collectFromScalar(e.getKey(), out);
169+
collectFromScalar(e.getValue(), out);
170+
}
171+
// lambdaExpr is traversed below via getChildren() (== [lambdaExpr]); don't recurse it here.
172+
}
173+
for (ScalarOperator child : s.getChildren()) {
174+
collectFromScalar(child, out);
175+
}
176+
}
177+
}

0 commit comments

Comments
 (0)