[BugFix] Fix unexpected Hint when UPDATE Stmt has CTE (StarRocks#51458)

Signed-off-by: stdpain <[email protected]>
vivo · Oct 2, 2024 · ecd7cd9 · ecd7cd9
1 parent e9eedcd
commit ecd7cd9
Show file tree

Hide file tree

Showing 4 changed files with 123 additions and 4 deletions.
diff --git a/be/src/exprs/agg/percentile_approx.h b/be/src/exprs/agg/percentile_approx.h
@@ -50,7 +50,11 @@ class PercentileApproxAggregateFunction final
             column_value = down_cast<const DoubleColumn*>(columns[0])->get_data()[row_num];
         }
 
-        DCHECK(!columns[1]->only_null());
+        if (columns[1]->only_null()) {
+            ctx->set_error("For percentile_approx the second argument is expected to be non-null.", false);
+            return;
+        }
+
         DCHECK(!columns[1]->is_null(0));
 
         int64_t prev_memory = data(state).percentile->mem_usage();

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/parser/HintCollector.java b/fe/fe-core/src/main/java/com/starrocks/sql/parser/HintCollector.java
@@ -101,7 +101,7 @@ public Void visitInsertStatement(StarRocksParser.InsertStatementContext context)
 
     @Override
     public Void visitUpdateStatement(StarRocksParser.UpdateStatementContext context) {
-        extractHintToRight(context);
+        extractHintToRight(context, context.UPDATE().getSymbol().getTokenIndex());
         if (!(context.fromClause() instanceof StarRocksParser.DualContext)) {
             StarRocksParser.FromContext fromContext = (StarRocksParser.FromContext) context.fromClause();
             if (fromContext.relations() != null) {
@@ -113,7 +113,7 @@ public Void visitUpdateStatement(StarRocksParser.UpdateStatementContext context)
 
     @Override
     public Void visitDeleteStatement(StarRocksParser.DeleteStatementContext context) {
-        extractHintToRight(context);
+        extractHintToRight(context, context.DELETE().getSymbol().getTokenIndex());
         if (context.using != null) {
             context.using.relation().stream().forEach(this::visit);
         }
@@ -196,4 +196,11 @@ private void extractHintToRight(ParserRuleContext ctx) {
             contextWithTokenMap.computeIfAbsent(ctx, e -> new ArrayList<>()).addAll(hintTokens);
         }
     }
+
+    private void extractHintToRight(ParserRuleContext ctx, int fromIndex) {
+        List<Token> hintTokens = tokenStream.getHiddenTokensToRight(fromIndex, HINT_CHANNEL);
+        if (hintTokens != null) {
+            contextWithTokenMap.computeIfAbsent(ctx, e -> new ArrayList<>()).addAll(hintTokens);
+        }
+    }
 }
diff --git a/test/sql/test_user_variables/R/test_user_variable b/test/sql/test_user_variables/R/test_user_variable
@@ -46,4 +46,67 @@ with tx as (select @var2 as x)
 select /*+ SET_USER_VARIABLE(@var2 = (select count(*) from (select l.c0 from t0 l join t0 r on l.c0 = r.c0 ) tb)) */ * from tx;
 -- result:
 40960
+-- !result
+CREATE TABLE `always1` (
+  `c0` int(11) NULL COMMENT ""
+) ENGINE=OLAP
+DUPLICATE KEY(`c0`)
+COMMENT "OLAP"
+DISTRIBUTED BY HASH(`c0`) BUCKETS 1
+PROPERTIES (
+"replication_num" = "1"
+);
+-- result:
+-- !result
+insert into always1 values(1);
+-- result:
+-- !result
+CREATE TABLE `alwaysnull` (
+  `c0` int(11) NULL COMMENT ""
+) ENGINE=OLAP
+DUPLICATE KEY(`c0`)
+COMMENT "OLAP"
+DISTRIBUTED BY HASH(`c0`) BUCKETS 1
+PROPERTIES (
+"replication_num" = "1"
+);
+-- result:
+-- !result
+insert into alwaysnull values(null);
+-- result:
+-- !result
+select      /*+ set_user_variable(@a = (select c0 * 0.1 from always1)) */ percentile_approx(c0, @a) is not null  from t0;
+-- result:
+1
+-- !result
+select      /*+ set_user_variable(@a = (select c0 * 0.1 from alwaysnull)) */ percentile_approx(c0, @a)  from t0;
+-- result:
+E: (1064, "Getting analyzing error. Detail message: percentile_approx requires the second parameter's type is numeric type.")
+-- !result
+select      percentile_approx(c0, cast(null as double))  from t0;
+-- result:
+[REGEX] .*percentile_approx the second argument is expected to be non-null.*
+-- !result
+CREATE TABLE `pk1` (
+  `c0` int(11) COMMENT "",
+  `dt` int(11) COMMENT ""
+) ENGINE=OLAP
+PRIMARY KEY(`c0`)
+COMMENT "OLAP"
+DISTRIBUTED BY HASH(`c0`) BUCKETS 1
+PROPERTIES (
+"replication_num" = "1"
+);
+-- result:
+-- !result
+insert into pk1 values (1,1);
+-- result:
+-- !result
+with cte as (select* from pk1)
+update /*+ set_user_variable(@a = (select 2)) */ pk1 set dt = @a where c0 = 1;
+-- result:
+-- !result
+select * from pk1;
+-- result:
+1	2
 -- !result
diff --git a/test/sql/test_user_variables/T/test_user_variable b/test/sql/test_user_variables/T/test_user_variable
@@ -29,4 +29,49 @@ select @var2;
 
 
 with tx as (select @var2 as x)
-select /*+ SET_USER_VARIABLE(@var2 = (select count(*) from (select l.c0 from t0 l join t0 r on l.c0 = r.c0 ) tb)) */ * from tx;
+select /*+ SET_USER_VARIABLE(@var2 = (select count(*) from (select l.c0 from t0 l join t0 r on l.c0 = r.c0 ) tb)) */ * from tx;
+
+CREATE TABLE `always1` (
+  `c0` int(11) NULL COMMENT ""
+) ENGINE=OLAP
+DUPLICATE KEY(`c0`)
+COMMENT "OLAP"
+DISTRIBUTED BY HASH(`c0`) BUCKETS 1
+PROPERTIES (
+"replication_num" = "1"
+);
+
+insert into always1 values(1);
+
+CREATE TABLE `alwaysnull` (
+  `c0` int(11) NULL COMMENT ""
+) ENGINE=OLAP
+DUPLICATE KEY(`c0`)
+COMMENT "OLAP"
+DISTRIBUTED BY HASH(`c0`) BUCKETS 1
+PROPERTIES (
+"replication_num" = "1"
+);
+
+insert into alwaysnull values(null);
+
+select      /*+ set_user_variable(@a = (select c0 * 0.1 from always1)) */ percentile_approx(c0, @a) is not null  from t0;
+select      /*+ set_user_variable(@a = (select c0 * 0.1 from alwaysnull)) */ percentile_approx(c0, @a)  from t0;
+select      percentile_approx(c0, cast(null as double))  from t0;
+
+
+CREATE TABLE `pk1` (
+  `c0` int(11) COMMENT "",
+  `dt` int(11) COMMENT ""
+) ENGINE=OLAP
+PRIMARY KEY(`c0`)
+COMMENT "OLAP"
+DISTRIBUTED BY HASH(`c0`) BUCKETS 1
+PROPERTIES (
+"replication_num" = "1"
+);
+insert into pk1 values (1,1);
+
+with cte as (select* from pk1)
+update /*+ set_user_variable(@a = (select 2)) */ pk1 set dt = @a where c0 = 1;
+select * from pk1;