Skip to content

Commit

Permalink
planner: fix incorrect maintenance of handleColHelper for recursive…
Browse files Browse the repository at this point in the history 
… CTE (#55732) (#55899)

close #55666
  • Loading branch information
@ti-chi-bot
ti-chi-bot authored Sep 6, 2024
1 parent 5315279 commit d00c55f
Show file tree
Hide file tree
Showing 4 changed files with 163 additions and 6 deletions.
34 changes: 29 additions & 5 deletions pkg/planner/core/logical_plan_builder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4547,6 +4547,9 @@ func getLatestVersionFromStatsTable(ctx sessionctx.Context, tblInfo *model.Table
return version
}

// tryBuildCTE considers the input tn as a reference to a CTE and tries to build the logical plan for it like building
// DataSource for normal tables.
// tryBuildCTE will push an entry into handleHelper when successful.
func (b *PlanBuilder) tryBuildCTE(ctx context.Context, tn *ast.TableName, asName *model.CIStr) (LogicalPlan, error) {
for i := len(b.outerCTEs) - 1; i >= 0; i-- {
cte := b.outerCTEs[i]
Expand All @@ -4571,6 +4574,7 @@ func (b *PlanBuilder) tryBuildCTE(ctx context.Context, tn *ast.TableName, asName
p := LogicalCTETable{name: cte.def.Name.String(), idForStorage: cte.storageID, seedStat: cte.seedStat, seedSchema: cte.seedLP.Schema()}.Init(b.ctx, b.getSelectOffset())
p.SetSchema(getResultCTESchema(cte.seedLP.Schema(), b.ctx.GetSessionVars()))
p.SetOutputNames(cte.seedLP.OutputNames())
b.handleHelper.pushMap(nil)
return p, nil
}

Expand Down Expand Up @@ -7390,6 +7394,8 @@ func isJoinHintSupportedInMPPMode(preferJoinType uint) bool {
return onesCount < 1
}

// buildCte prepares for a CTE. It works together with buildWith().
// It will push one entry into b.handleHelper.
func (b *PlanBuilder) buildCte(ctx context.Context, cte *ast.CommonTableExpression, isRecursive bool) (p LogicalPlan, err error) {
saveBuildingCTE := b.buildingCTE
b.buildingCTE = true
Expand Down Expand Up @@ -7425,6 +7431,7 @@ func (b *PlanBuilder) buildCte(ctx context.Context, cte *ast.CommonTableExpressi
}

// buildRecursiveCTE handles the with clause `with recursive xxx as xx`.
// It will push one entry into b.handleHelper.
func (b *PlanBuilder) buildRecursiveCTE(ctx context.Context, cte ast.ResultSetNode) error {
b.isCTE = true
cInfo := b.outerCTEs[len(b.outerCTEs)-1]
Expand Down Expand Up @@ -7454,6 +7461,7 @@ func (b *PlanBuilder) buildRecursiveCTE(ctx context.Context, cte ast.ResultSetNo
for i := 0; i < len(x.SelectList.Selects); i++ {
var p LogicalPlan
var err error
originalLen := b.handleHelper.stackTail

var afterOpr *ast.SetOprType
switch y := x.SelectList.Selects[i].(type) {
Expand All @@ -7465,6 +7473,22 @@ func (b *PlanBuilder) buildRecursiveCTE(ctx context.Context, cte ast.ResultSetNo
afterOpr = y.AfterSetOperator
}

// This is for maintain b.handleHelper instead of normal error handling. Since one error is expected if
// expectSeed && cInfo.useRecursive, error handling is in the "if expectSeed" block below.
if err == nil {
b.handleHelper.popMap()
} else {
// Be careful with this tricky case. One error is expected here when building the first recursive
// part, however, the b.handleHelper won't be restored if error occurs, which means there could be
// more than one entry pushed into b.handleHelper without being poped.
// For example: with recursive cte1 as (select ... union all select ... from tbl join cte1 ...) ...
// This violates the semantic of buildSelect() and buildSetOpr(), which should only push exactly one
// entry into b.handleHelper. So we use a special logic to restore the b.handleHelper here.
for b.handleHelper.stackTail > originalLen {
b.handleHelper.popMap()
}
}

if expectSeed {
if cInfo.useRecursive {
// 3. If it fail to build a plan, it may be the recursive part. Then we build the seed part plan, and rebuild it.
Expand Down Expand Up @@ -7495,14 +7519,11 @@ func (b *PlanBuilder) buildRecursiveCTE(ctx context.Context, cte ast.ResultSetNo
// Build seed part plan.
saveSelect := x.SelectList.Selects
x.SelectList.Selects = x.SelectList.Selects[:i]
// We're rebuilding the seed part, so we pop the result we built previously.
for _i := 0; _i < i; _i++ {
b.handleHelper.popMap()
}
p, err = b.buildSetOpr(ctx, x)
if err != nil {
return err
}
b.handleHelper.popMap()
x.SelectList.Selects = saveSelect
p, err = b.adjustCTEPlanOutputName(p, cInfo.def)
if err != nil {
Expand Down Expand Up @@ -7571,6 +7592,7 @@ func (b *PlanBuilder) buildRecursiveCTE(ctx context.Context, cte ast.ResultSetNo
limit.SetChildren(limit.Children()[:0]...)
cInfo.limitLP = limit
}
b.handleHelper.pushMap(nil)
return nil
default:
p, err := b.buildResultSetNode(ctx, x, true)
Expand Down Expand Up @@ -7666,7 +7688,9 @@ func (b *PlanBuilder) buildWith(ctx context.Context, w *ast.WithClause) ([]*cteI
b.outerCTEs[len(b.outerCTEs)-1].optFlag = b.optFlag
b.outerCTEs[len(b.outerCTEs)-1].isBuilding = false
b.optFlag = saveFlag
// each cte (select statement) will generate a handle map, pop it out here.
// buildCte() will push one entry into handleHelper. As said in comments for b.handleHelper, building CTE
// should not affect the handleColHelper, so we pop it out here, then buildWith() as a whole will not modify
// the handleColHelper.
b.handleHelper.popMap()
ctes = append(ctes, b.outerCTEs[len(b.outerCTEs)-1])
}
Expand Down
3 changes: 2 additions & 1 deletion pkg/planner/core/planbuilder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,8 @@ type PlanBuilder struct {
// If it's a aggregation, we pop the map and push a nil map since no handle information left.
// If it's a union, we pop all children's and push a nil map.
// If it's a join, we pop its children's out then merge them and push the new map to stack.
// If we meet a subquery, it's clearly that it's a independent problem so we just pop one map out when we finish building the subquery.
// If we meet a subquery or CTE, it's clearly that it's an independent problem so we just pop one map out when we
// finish building the subquery or CTE.
handleHelper *handleColHelper

hintProcessor *hint.QBHintHandler
Expand Down
64 changes: 64 additions & 0 deletions tests/integrationtest/r/cte.result
View file
Original file line number Diff line number Diff line change
Expand Up @@ -788,3 +788,67 @@ with cte1 as (select 1), cte2 as (select 2) select * from cte1 union (with cte2
1
1
3
drop table if exists t1, t2;
create table t1(a int, b int);
create table t2(a int, b int);
insert into t1 value(5,5);
insert into t2 value(1,1);
with recursive cte1 as (select 1 as a union all select cte1.a+1 from t1 join cte1 on t1.a > cte1.a) select * from cte1;
a
1
2
3
4
5
update t2 set b=2 where a in (with recursive cte1 as (select 1 as a union all select cte1.a+1 from t1 join cte1 on t1.a > cte1.a) select * from cte1);
select * from t2;
a b
1 2
delete from t2 where a in (with recursive cte1 as (select 1 as a union all select cte1.a+1 from t1 join cte1 on t1.a > cte1.a) select * from cte1);
select * from t2;
a b
drop table if exists table_abc1;
drop table if exists table_abc2;
drop table if exists table_abc3;
drop table if exists table_abc4;
CREATE TABLE `table_abc1` (
`column_abc1` varchar(10) DEFAULT NULL,
`column_abc2` varchar(10) DEFAULT NULL,
`column_abc3` varchar(10) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;
CREATE TABLE `table_abc3` (
`column_abc5` varchar(10) DEFAULT NULL,
`column_abc6` varchar(10) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;
CREATE TABLE `table_abc4` (
`column_abc3` varchar(10) DEFAULT NULL,
`column_abc7` varchar(10) DEFAULT NULL,
`column_abc5` varchar(10) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;
INSERT INTO `table_abc1` VALUES ('KTL157','KTL157','KTL157');
INSERT INTO `table_abc3` VALUES ('1000','20240819');
INSERT INTO `table_abc4` VALUES ('KTL157','test','1000');
DELETE FROM table_abc3 t_abc3
WHERE t_abc3.column_abc5 IN (
SELECT a.column_abc5
FROM (
WITH tree_cte1 AS (
WITH RECURSIVE tree_cte AS (
SELECT t.column_abc1, t.column_abc2, t.column_abc3, 0 AS lv
FROM table_abc1 t
WHERE t.column_abc1 IN ('KTL157', 'KTL159')
UNION ALL
SELECT t.column_abc1, t.column_abc2, t.column_abc3, tcte.lv + 1
FROM table_abc1 t
JOIN tree_cte tcte ON t.column_abc1 = tcte.column_abc2
WHERE tcte.lv <= 1
)
SELECT * FROM tree_cte
)
SELECT e.column_abc5
FROM (
SELECT DISTINCT * FROM tree_cte1
) aa
LEFT JOIN table_abc4 e ON e.column_abc3 = aa.column_abc3
) a
);
68 changes: 68 additions & 0 deletions tests/integrationtest/t/cte.test
View file
Original file line number Diff line number Diff line change
Expand Up @@ -336,3 +336,71 @@ INSERT INTO `t_dnmxh` VALUES (104,571000,NULL),(104,572000,44.37),(104,573000,59
WITH cte_0 AS (select distinct ref_0.wkey as c0, ref_0.pkey as c1, ref_0.c_xhsndb as c2 from t_dnmxh as ref_0 where (1 <= ( select ref_1.pkey not in ( select ref_5.wkey as c0 from t_dnmxh as ref_5 where (ref_5.wkey < ( select ref_6.pkey as c0 from t_cqmg3b as ref_6 where 88 between 96 and 76)) ) as c0 from (t_cqmg3b as ref_1 left outer join t_dnmxh as ref_2 on (ref_1.wkey = ref_2.wkey )) where ref_0.c_xhsndb is NULL union select 33 <= 91 as c0 from t_cqmg3b as ref_8 ))), cte_1 AS (select ref_9.wkey as c0, ref_9.pkey as c1, ref_9.c_anpf_c as c2, ref_9.c_b_fp_c as c3, ref_9.c_ndccfb as c4, ref_9.c_8rswc as c5 from t_cqmg3b as ref_9) select count(1) from cte_0 as ref_10 where case when 56 < 50 then case when 100 in ( select distinct ref_11.c4 as c0 from cte_1 as ref_11 where (ref_11.c4 > ( select ref_13.pkey as c0 from t_dnmxh as ref_13 where (ref_13.wkey > ( select distinct ref_11.c1 as c0 from cte_0 as ref_14)) )) or (1 = 1)) then null else null end else '7mxv6' end not like 'ki4%vc';
#case
with cte1 as (select 1), cte2 as (select 2) select * from cte1 union (with cte2 as (select 3) select * from cte2 union all select * from cte2) order by 1;

# Tests for PR #55732
drop table if exists t1, t2;
create table t1(a int, b int);
create table t2(a int, b int);
insert into t1 value(5,5);
insert into t2 value(1,1);
with recursive cte1 as (select 1 as a union all select cte1.a+1 from t1 join cte1 on t1.a > cte1.a) select * from cte1;
# This UPDATE should update t2.b to 2
update t2 set b=2 where a in (with recursive cte1 as (select 1 as a union all select cte1.a+1 from t1 join cte1 on t1.a > cte1.a) select * from cte1);
select * from t2;
# This DELETE should delete all rows in t2
delete from t2 where a in (with recursive cte1 as (select 1 as a union all select cte1.a+1 from t1 join cte1 on t1.a > cte1.a) select * from cte1);
select * from t2;

# Issue #55666
drop table if exists table_abc1;
drop table if exists table_abc2;
drop table if exists table_abc3;
drop table if exists table_abc4;

CREATE TABLE `table_abc1` (
`column_abc1` varchar(10) DEFAULT NULL,
`column_abc2` varchar(10) DEFAULT NULL,
`column_abc3` varchar(10) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;

CREATE TABLE `table_abc3` (
`column_abc5` varchar(10) DEFAULT NULL,
`column_abc6` varchar(10) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;

CREATE TABLE `table_abc4` (
`column_abc3` varchar(10) DEFAULT NULL,
`column_abc7` varchar(10) DEFAULT NULL,
`column_abc5` varchar(10) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;

INSERT INTO `table_abc1` VALUES ('KTL157','KTL157','KTL157');
INSERT INTO `table_abc3` VALUES ('1000','20240819');
INSERT INTO `table_abc4` VALUES ('KTL157','test','1000');

DELETE FROM table_abc3 t_abc3
WHERE t_abc3.column_abc5 IN (
SELECT a.column_abc5
FROM (
WITH tree_cte1 AS (
WITH RECURSIVE tree_cte AS (
SELECT t.column_abc1, t.column_abc2, t.column_abc3, 0 AS lv
FROM table_abc1 t
WHERE t.column_abc1 IN ('KTL157', 'KTL159')
UNION ALL
SELECT t.column_abc1, t.column_abc2, t.column_abc3, tcte.lv + 1
FROM table_abc1 t
JOIN tree_cte tcte ON t.column_abc1 = tcte.column_abc2
WHERE tcte.lv <= 1
)
SELECT * FROM tree_cte
)
SELECT e.column_abc5
FROM (
SELECT DISTINCT * FROM tree_cte1
) aa
LEFT JOIN table_abc4 e ON e.column_abc3 = aa.column_abc3
) a
);


0 comments on commit d00c55f

Please sign in to comment.