imp: merge union to aggregate plan (#180)

pkg/runtime/optimize/optimizer.go

@@ -315,13 +315,17 @@ func (o optimizer) optimizeSelect(ctx context.Context, conn proto.VConn, stmt *r
 	}
 
 	unionPlan := &plan.UnionPlan{
-		Plans:      plans,
+		Plans: plans,
+	}
+
+	// TODO: order/groupBy/aggregate
+	aggregate := &plan.AggregatePlan{
+		UnionPlan:  unionPlan,
 		Combiner:   transformer.NewCombinerManager(),
 		AggrLoader: transformer.LoadAggrs(stmt.Select),
 	}
-	// TODO: order/groupBy/aggregate
 
-	return unionPlan, nil
+	return aggregate, nil
 }
 
 func (o optimizer) optimizeUpdate(ctx context.Context, conn proto.VConn, stmt *rast.UpdateStatement, args []interface{}) (proto.Plan, error) {

pkg/runtime/plan/aggregate.go

@@ -0,0 +1,50 @@
+/*
+ * 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 plan
+
+import (
+	"context"
+)
+
+import (
+	"github.com/pkg/errors"
+)
+
+import (
+	"github.com/arana-db/arana/pkg/proto"
+	"github.com/arana-db/arana/pkg/transformer"
+)
+
+type AggregatePlan struct {
+	transformer.Combiner
+	AggrLoader *transformer.AggrLoader
+	UnionPlan  *UnionPlan
+}
+
+func (a *AggregatePlan) Type() proto.PlanType {
+	return proto.PlanTypeQuery
+}
+
+func (a *AggregatePlan) ExecIn(ctx context.Context, conn proto.VConn) (proto.Result, error) {
+	res, err := a.UnionPlan.ExecIn(ctx, conn)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	return a.Combiner.Merge(res, a.AggrLoader)
+}

pkg/runtime/plan/union.go

@@ -27,14 +27,11 @@ import (
 
 import (
 	"github.com/arana-db/arana/pkg/proto"
-	"github.com/arana-db/arana/pkg/transformer"
 )
 
 // UnionPlan merges multiple query plan.
 type UnionPlan struct {
 	Plans []proto.Plan
-	transformer.Combiner
-	AggrLoader *transformer.AggrLoader
 }
 
 func (u UnionPlan) Type() proto.PlanType {
@@ -51,7 +48,7 @@ func (u UnionPlan) ExecIn(ctx context.Context, conn proto.VConn) (proto.Result,
 		results = append(results, res)
 	}
 
-	return u.Combiner.Merge(compositeResult(results), u.AggrLoader)
+	return compositeResult(results), nil
 }
 
 type compositeResult []proto.Result