fix limit

Author: xitongsys

Diff for: eplan/eplan.go

@@ -335,13 +335,15 @@ func createEPlan(node PlanNode, ePlanNodes *[]ENode, executorHeap *util.Heap, pn
 		if err != nil {
 			return res, err
 		}
+		inputs := []pb.Location{}
 		for _, inputNode := range inputNodes {
-			for _, input := range inputNode.GetOutputs() {
-				output := executorHeap.GetExecutorLoc()
-				output.ChannelIndex = 0
-				res = append(res, NewEPlanLimitNode(nodea, input, output))
-			}
+			inputs = append(inputs, inputNode.GetOutputs()...)
 		}
+
+		limitNodeLoc := executorHeap.GetExecutorLoc()
+		limitNodeLoc.ChannelIndex = 0
+		res = append(res, NewEPlanLimitNode(nodea, inputs, limitNodeLoc))
+
 		*ePlanNodes = append(*ePlanNodes, res...)
 		return res, nil
 

Diff for: eplan/limit.go

@@ -8,7 +8,7 @@ import (
 
 type EPlanLimitNode struct {
 	Location    pb.Location
-	Input       pb.Location
+	Inputs      []pb.Location
 	Output      pb.Location
 	LimitNumber *int64
 	Metadata    *metadata.Metadata
@@ -19,7 +19,7 @@ func (self *EPlanLimitNode) GetNodeType() EPlanNodeType {
 }
 
 func (self *EPlanLimitNode) GetInputs() []pb.Location {
-	return []pb.Location{self.Input}
+	return self.Inputs
 }
 
 func (self *EPlanLimitNode) GetOutputs() []pb.Location {
@@ -30,10 +30,10 @@ func (self *EPlanLimitNode) GetLocation() pb.Location {
 	return self.Location
 }
 
-func NewEPlanLimitNode(node *PlanLimitNode, input, output pb.Location) *EPlanLimitNode {
+func NewEPlanLimitNode(node *PlanLimitNode, inputs []pb.Location, output pb.Location) *EPlanLimitNode {
 	return &EPlanLimitNode{
 		Location:    output,
-		Input:       input,
+		Inputs:      inputs,
 		Output:      output,
 		LimitNumber: node.LimitNumber,
 		Metadata:    node.GetMetadata(),

Diff for: executor/limit.go

@@ -23,7 +23,10 @@ func (self *Executor) SetInstructionLimit(instruction *pb.Instruction) (err erro
 	}
 	self.Instruction = instruction
 	self.EPlanNode = &enode
-	self.InputLocations = []*pb.Location{&enode.Input}
+	self.InputLocations = []*pb.Location{}
+	for i := 0; i < len(enode.Inputs); i++ {
+		self.InputLocations = append(self.InputLocations, &(enode.Inputs[i]))
+	}
 	self.OutputLocations = []*pb.Location{&enode.Output}
 	return nil
 }
@@ -79,9 +82,24 @@ func (self *Executor) RunLimit() (err error) {
 			if err != nil {
 				return err
 			}
-			readRowCnt += int64(rg.GetRowsNumber())
-			if err = rbWriter.Write(rg); err != nil {
-				return err
+			if readRowCnt+int64(rg.GetRowsNumber()) <= *(enode.LimitNumber) {
+				readRowCnt += int64(rg.GetRowsNumber())
+				if err = rbWriter.Write(rg); err != nil {
+					return err
+				}
+
+			} else {
+				for readRowCnt < *(enode.LimitNumber) {
+					row, err := rg.Read()
+					if err != nil {
+						return err
+					}
+					if err = rbWriter.WriteRow(row); err != nil {
+						return err
+					}
+					readRowCnt++
+
+				}
 			}
 		}
 	}

Diff for: master/ui_task_info.go

@@ -116,6 +116,7 @@ func DrawNode(canvas *svg.SVG, node *SVGNode) {
 		"DISTINCT LOCAL":        "stroke-width:0; fill:rgb(15,18,98);",
 		"DISTINCT GLOBAL":       "stroke-width:0; fill:rgb(15,58,148);",
 		"SHOW":                  "stroke-width:0; fill:rgb(125,148,148);",
+		"BALANCE":               "stroke-width:0; fill:rgb(125,18,148);",
 		"UNKNOWN":               "stroke-width:0; fill:rgb(181,181,181);",
 	}
 	style := NodeStyle[node.NodeType]