feat: implement complex select with orderby (#379)

integration_test/scene/db_tbl_rw/integration_test.go

@@ -55,8 +55,8 @@ func (s *IntegrationSuite) TestDBTBLRWScene() {
 		db = s.DB()
 		t  = s.T()
 	)
-	//tx, err := db.Begin()
-	//assert.NoError(t, err, "should begin a new tx")
+	// tx, err := db.Begin()
+	// assert.NoError(t, err, "should begin a new tx")
 
 	cases := s.TestCases()
 	for _, sqlCase := range cases.ExecCases {

pkg/mysql/fields.go

@@ -62,6 +62,15 @@ type Field struct {
 	defaultValue       []byte
 }
 
+func (mf *Field) SetOrgName(name string) {
+	mf.orgName = name
+}
+
+func (mf *Field) SetName(name string) {
+	// TODO: should resize columnLength???
+	mf.name = name
+}
+
 func (mf *Field) FieldType() mysql.FieldType {
 	return mf.fieldType
 }

pkg/runtime/ast/ast.go

@@ -895,6 +895,10 @@ func (cc *convCtx) convFieldList(node *ast.FieldList) []SelectElement {
 			continue
 		}
 
+		getOriginalText := func() string {
+			return field.Text()
+		}
+
 		alias := field.AsName.String()
 		switch t := cc.convExpr(field.Expr).(type) {
 		case *AtomPredicateNode:
@@ -906,33 +910,39 @@ func (cc *convCtx) convFieldList(node *ast.FieldList) []SelectElement {
 				})
 			case *FunctionCallExpressionAtom:
 				ret = append(ret, &SelectElementFunction{
-					inner: a.F,
-					alias: alias,
+					inner:        a.F,
+					alias:        alias,
+					originalText: getOriginalText(),
 				})
 			case *ConstantExpressionAtom:
 				ret = append(ret, &SelectElementExpr{
-					inner: exprAtomToNode(a),
-					alias: alias,
+					inner:        exprAtomToNode(a),
+					alias:        alias,
+					originalText: getOriginalText(),
 				})
 			case *MathExpressionAtom:
 				ret = append(ret, &SelectElementExpr{
-					inner: exprAtomToNode(a),
-					alias: alias,
+					inner:        exprAtomToNode(a),
+					alias:        alias,
+					originalText: getOriginalText(),
 				})
 			case *UnaryExpressionAtom:
 				ret = append(ret, &SelectElementExpr{
-					inner: exprAtomToNode(a),
-					alias: alias,
+					inner:        exprAtomToNode(a),
+					alias:        alias,
+					originalText: getOriginalText(),
 				})
 			case *NestedExpressionAtom:
 				ret = append(ret, &SelectElementExpr{
-					inner: exprAtomToNode(a),
-					alias: alias,
+					inner:        exprAtomToNode(a),
+					alias:        alias,
+					originalText: getOriginalText(),
 				})
 			case *SystemVariableExpressionAtom:
 				ret = append(ret, &SelectElementExpr{
-					inner: exprAtomToNode(a),
-					alias: alias,
+					inner:        exprAtomToNode(a),
+					alias:        alias,
+					originalText: getOriginalText(),
 				})
 			default:
 				panic(fmt.Sprintf("todo: unsupported select element type %T!", a))

pkg/runtime/ast/expression_atom.go

@@ -249,6 +249,10 @@ func (c *ConstantExpressionAtom) CntParams() int {
 
 type ColumnNameExpressionAtom []string
 
+func NewSingleColumnNameExpressionAtom(name string) ColumnNameExpressionAtom {
+	return ColumnNameExpressionAtom{name}
+}
+
 func (c ColumnNameExpressionAtom) Accept(visitor Visitor) (interface{}, error) {
 	return visitor.VisitAtomColumn(c)
 }

pkg/runtime/ast/select_element.go

@@ -38,17 +38,27 @@ type selectElementPhantom struct{}
 type SelectElement interface {
 	Node
 	Restorer
+
 	// Alias returns the alias if available.
 	Alias() string
+
 	// ToSelectString converts current SelectElement to SQL SELECT string and return it.
 	ToSelectString() string
+
+	// DisplayName returns the actual display name of field.
+	DisplayName() string
+
 	phantom() selectElementPhantom
 }
 
 type SelectElementAll struct {
 	prefix string
 }
 
+func (s *SelectElementAll) DisplayName() string {
+	return s.ToSelectString()
+}
+
 func (s *SelectElementAll) Accept(visitor Visitor) (interface{}, error) {
 	return visitor.VisitSelectElementWildcard(s)
 }
@@ -85,8 +95,19 @@ func (s *SelectElementAll) phantom() selectElementPhantom {
 }
 
 type SelectElementExpr struct {
-	inner ExpressionNode
-	alias string
+	inner        ExpressionNode
+	alias        string
+	originalText string
+}
+
+func (s *SelectElementExpr) DisplayName() string {
+	if len(s.alias) > 0 {
+		return s.alias
+	}
+	if len(s.originalText) > 0 {
+		return s.originalText
+	}
+	return s.ToSelectString()
 }
 
 func (s *SelectElementExpr) Accept(visitor Visitor) (interface{}, error) {
@@ -123,8 +144,19 @@ func (s *SelectElementExpr) phantom() selectElementPhantom {
 }
 
 type SelectElementFunction struct {
-	inner Node // *Function or *AggrFunction
-	alias string
+	inner        Node // *Function or *AggrFunction
+	alias        string
+	originalText string
+}
+
+func (s *SelectElementFunction) DisplayName() string {
+	if len(s.alias) > 0 {
+		return s.alias
+	}
+	if len(s.originalText) > 0 {
+		return s.originalText
+	}
+	return s.ToSelectString()
 }
 
 func (s *SelectElementFunction) Accept(visitor Visitor) (interface{}, error) {
@@ -217,6 +249,13 @@ type SelectElementColumn struct {
 	alias string
 }
 
+func (s *SelectElementColumn) DisplayName() string {
+	if len(s.alias) > 0 {
+		return s.alias
+	}
+	return s.Suffix()
+}
+
 func (s *SelectElementColumn) Accept(visitor Visitor) (interface{}, error) {
 	return visitor.VisitSelectElementColumn(s)
 }

pkg/runtime/optimize/dml/aggregate_visitor.go

@@ -48,12 +48,17 @@ func (av *aggregateVisitor) VisitSelectStatement(node *ast.SelectStatement) (int
 			return nil, errors.WithStack(err)
 		}
 
-		switch next.(type) {
+		switch val := next.(type) {
 		case *ext.WeakSelectElement:
 			rebuilds = append(rebuilds, &ext.WeakSelectElement{
 				SelectElement: res.(ast.SelectElement),
 			})
 			av.hasWeak = true
+		case *ext.WeakAliasSelectElement:
+			rebuilds = append(rebuilds, &ext.WeakAliasSelectElement{
+				SelectElement: res.(ast.SelectElement),
+				WeakAlias:     val.WeakAlias,
+			})
 		default:
 			rebuilds = append(rebuilds, res.(ast.SelectElement))
 		}

pkg/runtime/optimize/dml/ext/weak_alias.go

@@ -0,0 +1,55 @@
+/*
+ * 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 ext
+
+import (
+	"strings"
+)
+
+import (
+	"github.com/pkg/errors"
+)
+
+import (
+	"github.com/arana-db/arana/pkg/runtime/ast"
+)
+
+type WeakAliasSelectElement struct {
+	ast.SelectElement
+	WeakAlias string
+}
+
+func (re WeakAliasSelectElement) Prev() ast.SelectElement {
+	if p, ok := re.SelectElement.(SelectElementProvider); ok {
+		return p.Prev()
+	}
+	return re.SelectElement
+}
+
+func (re WeakAliasSelectElement) Restore(flag ast.RestoreFlag, sb *strings.Builder, args *[]int) error {
+	if err := re.SelectElement.Restore(flag, sb, args); err != nil {
+		return errors.WithStack(err)
+	}
+	sb.WriteString(" AS ")
+	ast.WriteID(sb, re.WeakAlias)
+	return nil
+}
+
+func (re WeakAliasSelectElement) Alias() string {
+	return re.WeakAlias
+}

pkg/runtime/optimize/dml/select.go

@@ -253,6 +253,17 @@ func optimizeSelect(ctx context.Context, o *optimize.Optimizer) (proto.Plan, err
 		}
 	}
 
+	// FIXME: tuning, avoid rename everytime.
+
+	// Rename return fields as normalized:
+	// For the query of "SELECT foo+1, avg(score) FROM xxx WHERE ...", will return columns:
+	//   BEFORE: | `foo`+1 | AVG(`score`) |
+	//      NOW: | foo+1 | avg(score) |
+	tmpPlan = &dml.RenamePlan{
+		Plan:       tmpPlan,
+		RenameList: analysis.normalizedFields,
+	}
+
 	return tmpPlan, nil
 }