[improvement](plsql) Select statement supports insert into variables a…

…pache#31574

fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4

@@ -137,6 +137,7 @@ BROKER: 'BROKER';
 BUCKETS: 'BUCKETS';
 BUILD: 'BUILD';
 BUILTIN: 'BUILTIN';
+BULK: 'BULK';
 BY: 'BY';
 CACHED: 'CACHED';
 CALL: 'CALL';
@@ -154,6 +155,7 @@ CLUSTER: 'CLUSTER';
 CLUSTERS: 'CLUSTERS';
 COLLATE: 'COLLATE';
 COLLATION: 'COLLATION';
+COLLECT: 'COLLECT';
 COLUMN: 'COLUMN';
 COLUMNS: 'COLUMNS';
 COMMENT: 'COMMENT';

fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4

@@ -301,6 +301,7 @@ queryPrimary
 
 querySpecification
     : selectClause
+      intoClause?
       fromClause?
       whereClause?
       aggClause?
@@ -337,6 +338,19 @@ fromClause
     : FROM relations
     ;
 
+// For PL-SQL
+intoClause
+    : bulkCollectClause? INTO (tableRow | identifier) (COMMA (tableRow | identifier))*
+    ;
+
+bulkCollectClause :
+       BULK COLLECT
+     ;
+
+tableRow :
+      identifier LEFT_PAREN INTEGER_VALUE RIGHT_PAREN
+    ;
+
 relations
     : relation (COMMA relation)*
     ;
@@ -947,6 +961,7 @@ nonReserved
     | BUCKETS
     | BUILD
     | BUILTIN
+    | BULK
     | CACHED
     | CALL
     | CATALOG
@@ -958,6 +973,7 @@ nonReserved
     | CLUSTER
     | CLUSTERS
     | COLLATION
+    | COLLECT
     | COLUMNS
     | COMMENT
     | COMMIT

fe/fe-core/src/main/antlr4/org/apache/doris/nereids/PLLexer.g4

@@ -38,7 +38,6 @@ BINARY_INTEGER: 'BINARY_INTEGER';
 BIT: 'BIT';
 BODY: 'BODY';
 BREAK: 'BREAK';
-BULK: 'BULK';
 BYTE: 'BYTE';
 CALLER: 'CALLER';
 CASCADE: 'CASCADE';
@@ -47,7 +46,6 @@ CLIENT: 'CLIENT';
 CLOSE: 'CLOSE';
 CLUSTERED: 'CLUSTERED';
 CMP: 'CMP';
-COLLECT: 'COLLECT';
 COLLECTION: 'COLLECTION';
 COMPRESS: 'COMPRESS';
 CONCAT: 'CONCAT';

fe/fe-core/src/main/antlr4/org/apache/doris/nereids/PLParser.g4

@@ -426,7 +426,7 @@ open_stmt :             // OPEN cursor statement
      ;
 
 fetch_stmt :            // FETCH cursor statement
-       FETCH FROM? ident_pl bulk_collect_clause? INTO ident_pl (COMMA ident_pl)* fetch_limit?
+       FETCH FROM? ident_pl bulkCollectClause? INTO ident_pl (COMMA ident_pl)* fetch_limit?
      ;
 
 fetch_limit:
@@ -521,10 +521,6 @@ using_clause :          // USING var,... clause
        USING expr (COMMA expr)*
      ;
 
-bulk_collect_clause :
-       BULK COLLECT
-     ;
-
 bool_expr :                               // Boolean condition
        NOT? LEFT_PAREN bool_expr RIGHT_PAREN
      | bool_expr bool_expr_logical_operator bool_expr
@@ -782,7 +778,6 @@ non_reserved_words :                      // Tokens that are not reserved words
      | BIT
      | BODY
      | BREAK
-     | BULK
      | BYTE
      | CALLER
      | CASCADE
@@ -791,7 +786,6 @@ non_reserved_words :                      // Tokens that are not reserved words
      | CLOSE
      | CLUSTERED
      | CMP
-     | COLLECT
      | COLLECTION
      | COMPRESS
      | CONSTANT

fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java

@@ -1226,6 +1226,9 @@ public LogicalPlan visitRegularQuerySpecification(RegularQuerySpecificationConte
             } else {
                 relation = visitFromClause(ctx.fromClause());
             }
+            if (ctx.intoClause() != null && !ConnectContext.get().isRunProcedure()) {
+                throw new ParseException("Only procedure supports insert into variables", selectCtx);
+            }
             selectPlan = withSelectQuerySpecification(
                     ctx, relation,
                     selectCtx,

fe/fe-core/src/main/java/org/apache/doris/plsql/Stmt.java

@@ -39,14 +39,19 @@
 import org.apache.doris.nereids.PLParser.If_plsql_stmtContext;
 import org.apache.doris.nereids.PLParser.If_tsql_stmtContext;
 import org.apache.doris.nereids.PLParser.Include_stmtContext;
+import org.apache.doris.nereids.PLParser.IntoClauseContext;
 import org.apache.doris.nereids.PLParser.Leave_stmtContext;
 import org.apache.doris.nereids.PLParser.Open_stmtContext;
 import org.apache.doris.nereids.PLParser.Print_stmtContext;
+import org.apache.doris.nereids.PLParser.QueryPrimaryDefaultContext;
 import org.apache.doris.nereids.PLParser.Quit_stmtContext;
+import org.apache.doris.nereids.PLParser.RegularQuerySpecificationContext;
 import org.apache.doris.nereids.PLParser.Resignal_stmtContext;
 import org.apache.doris.nereids.PLParser.Return_stmtContext;
 import org.apache.doris.nereids.PLParser.Set_current_schema_optionContext;
 import org.apache.doris.nereids.PLParser.Signal_stmtContext;
+import org.apache.doris.nereids.PLParser.StatementDefaultContext;
+import org.apache.doris.nereids.PLParser.TableRowContext;
 import org.apache.doris.nereids.PLParser.Unconditional_loop_stmtContext;
 import org.apache.doris.nereids.PLParser.Values_into_stmtContext;
 import org.apache.doris.nereids.PLParser.While_stmtContext;
@@ -81,7 +86,7 @@ public class Stmt {
 
     boolean trace = false;
     ResultListener resultListener = ResultListener.NONE;
-    private QueryExecutor queryExecutor;
+    private final QueryExecutor queryExecutor;
 
     Stmt(Exec e, QueryExecutor queryExecutor) {
         exec = e;
@@ -186,25 +191,55 @@ public Integer statement(ParserRuleContext ctx) {
         return 0;
     }
 
+    /**
+     * Get INTO clause
+     */
+    IntoClauseContext getIntoClause(ParserRuleContext ctx) {
+        if (ctx.getChild(0) instanceof StatementDefaultContext) {
+            ParserRuleContext queryTermDefaultCtx = ((StatementDefaultContext) ctx.getChild(0)).query().queryTerm();
+            if (queryTermDefaultCtx.getChild(0) instanceof QueryPrimaryDefaultContext) {
+                ParserRuleContext queryPrimaryDefaultContext
+                        = ((QueryPrimaryDefaultContext) queryTermDefaultCtx.getChild(0));
+                if (queryPrimaryDefaultContext.getChild(0) instanceof RegularQuerySpecificationContext) {
+                    return ((RegularQuerySpecificationContext) queryPrimaryDefaultContext.getChild(0)).intoClause();
+                }
+            }
+        }
+        return null;
+    }
+
     /**
      * Get number of elements in INTO or var=col assignment clause
      */
     int getIntoCount(ParserRuleContext ctx) {
-        // TODO
+        IntoClauseContext into = getIntoClause(ctx);
+        if (into != null) {
+            return into.identifier().size() + into.tableRow().size();
+        }
+        // TODO support var=col assignment clause
         return 0;
     }
 
     /**
      * Get variable name assigned in INTO or var=col clause by index
      */
     String getIntoVariable(ParserRuleContext ctx, int idx) {
-        // TODO
+        IntoClauseContext into = getIntoClause(ctx);
+        if (into != null) {
+            return into.tableRow(idx) != null ? into.tableRow(idx).identifier().getText()
+                    : into.identifier(idx).getText();
+        }
+        // TODO support var=col assignment clause
         return null;
     }
 
     private int getIntoTableIndex(ParserRuleContext ctx, int idx) {
-        // TODO
-        return 0;
+        IntoClauseContext into = getIntoClause(ctx);
+        TableRowContext row = into.tableRow(idx);
+        if (row == null) {
+            throw new RuntimeException("Missing into table index");
+        }
+        return Integer.parseInt(row.INTEGER_VALUE().getText());
     }
 
     private void populateVariable(ParserRuleContext ctx, QueryResult query, int columnIndex) throws AnalysisException {
@@ -380,7 +415,7 @@ public Integer fetch(Fetch_stmtContext ctx) {
             int cols = ctx.ident_pl().size() - 1;
             QueryResult queryResult = cursor.getQueryResult();
 
-            if (ctx.bulk_collect_clause() != null) {
+            if (ctx.bulkCollectClause() != null) {
                 long limit = ctx.fetch_limit() != null ? evalPop(ctx.fetch_limit().expr()).longValue() : -1;
                 long rowIndex = 1;
                 List<Table> tables = exec.intoTables(ctx, intoVariableNames(ctx, cols));
@@ -545,7 +580,7 @@ public Integer assignFromSelect(Assignment_stmt_select_itemContext ctx) {
                         if (trace) {
                             trace(ctx, "COLUMN: " + query.metadata().columnName(i) + ", " + query.metadata()
                                     .columnTypeName(i));
-                            trace(ctx, "SET " + var.getName() + " = " + var.toString());
+                            trace(ctx, "SET " + var.getName() + " = " + var);
                         }
                     } else if (trace) {
                         trace(ctx, "Variable not found: " + ctx.ident_pl(i).getText());