From d740a0cda7fbf63ce34092ad6cf9b209bab0470d Mon Sep 17 00:00:00 2001
From: nagamocha3000 <nagamocha@gmail.com>
Date: Tue, 23 Feb 2021 18:18:23 +0300
Subject: [PATCH] Add execution for median

---
 executor.go      | 76 ++++++++++++++++++++++++++++++++++++++++++++++++
 executor_test.go | 20 +++++--------
 2 files changed, 84 insertions(+), 12 deletions(-)

diff --git a/executor.go b/executor.go
index 85803be68..427691ed4 100644
--- a/executor.go
+++ b/executor.go
@@ -830,6 +830,9 @@ func (e *executor) executeCall(ctx context.Context, qcx *Qcx, index string, c *p
 	case "Limit":
 		res, err := e.executeLimitCall(ctx, qcx, index, c, shards, opt)
 		return res, errors.Wrapf(err, "executeLimitCall %v", shardSlice(shards))
+	case "Percentile":
+		res, err := e.executePercentile(ctx, qcx, index, c, shards, opt)
+		return res, errors.Wrapf(err, "executePercentile %v", shardSlice(shards))
 	default: // e.g. "Row", "Union", "Intersect" or anything that returns a bitmap.
 		statFn()
 		res, err := e.executeBitmapCall(ctx, qcx, index, c, shards, opt)
@@ -1293,6 +1296,79 @@ func (e *executor) executeMax(ctx context.Context, qcx *Qcx, index string, c *pq
 	return other, nil
 }
 
+// executePercentile executes a Percentile() call.
+func (e *executor) executePercentile(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (_ ValCount, err error) {
+	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executePercentile")
+	defer span.Finish()
+
+	if field := c.Args["field"]; field == "" {
+		return ValCount{}, errors.New("Percentile(): field required")
+	}
+	fieldName, _, _ := c.StringArg("field")
+
+	// get min
+	q, _ := pql.ParseString(fmt.Sprintf(`Min(field="%s")`, fieldName))
+	minCall := q.Calls[0]
+	minVal, err := e.executeMin(ctx, qcx, index, minCall, shards, opt)
+	if err != nil {
+		return ValCount{}, errors.Wrap(err, "executing Min call for Percentile")
+	}
+
+	// get max
+	q, _ = pql.ParseString(fmt.Sprintf(`Max(field="%s")`, fieldName))
+	maxCall := q.Calls[0]
+	maxVal, err := e.executeMax(ctx, qcx, index, maxCall, shards, opt)
+	if err != nil {
+		return ValCount{}, errors.Wrap(err, "executing Max call for Percentile")
+	}
+	// set up reusables
+	countQuery, _ := pql.ParseString("Count(Row(fld < 0))")
+	countCall := countQuery.Calls[0]
+	rangeQuery, _ := pql.ParseString("Row(fld < 0)")
+	rangeCall := rangeQuery.Calls[0]
+
+	min, max := minVal.Val, maxVal.Val
+	// estimate nth val, eg median when nth=0.5
+	for min < max {
+		possibleNthVal := (max - min) / 2
+		// get left count
+		rangeCall.Args[fieldName] = &pql.Condition{
+			Op:    pql.Token(pql.LT),
+			Value: possibleNthVal,
+		}
+		countCall.Children = []*pql.Call{rangeCall}
+		leftCountUint64, err := e.executeCount(ctx, qcx, index, countCall, shards, opt)
+		if err != nil {
+			return ValCount{}, errors.Wrap(err, "executing Count call L for Percentile")
+		}
+		leftCount := int64(leftCountUint64)
+
+		// get right count
+		rangeCall.Args[fieldName] = &pql.Condition{
+			Op:    pql.Token(pql.GT),
+			Value: possibleNthVal,
+		}
+		countCall.Children = []*pql.Call{rangeCall}
+		rightCountUint64, err := e.executeCount(ctx, qcx, index, countCall, shards, opt)
+		if err != nil {
+			return ValCount{}, errors.Wrap(err, "executing Count call R for Percentile")
+		}
+		rightCount := int64(rightCountUint64)
+
+		// binary search
+		if leftCount > rightCount {
+			max = possibleNthVal - 1
+		} else if leftCount < rightCount {
+			min = possibleNthVal + 1
+		} else {
+			return ValCount{Val: possibleNthVal, Count: 1}, nil
+		}
+	}
+
+	return ValCount{Val: min, Count: 1}, nil
+
+}
+
 // executeMinRow executes a MinRow() call.
 func (e *executor) executeMinRow(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (_ interface{}, err error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMinRow")
diff --git a/executor_test.go b/executor_test.go
index 605ec8bb4..80a7e2065 100644
--- a/executor_test.go
+++ b/executor_test.go
@@ -6871,13 +6871,13 @@ func variousQueriesOnPercentiles(t *testing.T, clusterSize int) {
 	// get binned to the same hour
 	c.CreateField(t, "users", pilosa.IndexOptions{Keys: true, TrackExistence: true}, "net_worth", pilosa.OptFieldTypeInt(-1000, 1000))
 	c.ImportIntKey(t, "users", "net_worth", []test.IntKey{
-		{Key: "user1", Val: 1},
-		{Key: "user2", Val: 2},
-		{Key: "user3", Val: 3},
-		{Key: "user4", Val: 4},
-		{Key: "user5", Val: 5},
-		{Key: "user6", Val: 6},
-		{Key: "user7", Val: 7},
+		{Key: "user1", Val: 10},
+		{Key: "user2", Val: 20},
+		{Key: "user3", Val: 30},
+		{Key: "user4", Val: 40},
+		{Key: "user5", Val: 50},
+		{Key: "user6", Val: 60},
+		{Key: "user7", Val: 70},
 	})
 
 	splitSortBackToCSV := func(csvStr string) string {
@@ -6886,10 +6886,6 @@ func variousQueriesOnPercentiles(t *testing.T, clusterSize int) {
 		return strings.Join(ss, "\n") + "\n"
 	}
 
-	toCSV := func(s string) string {
-		return strings.Join(strings.Split(s, " "), "\n") + "\n"
-	}
-
 	type testCase struct {
 		query       string
 		qrVerifier  func(t *testing.T, resp pilosa.QueryResponse)
@@ -6900,7 +6896,7 @@ func variousQueriesOnPercentiles(t *testing.T, clusterSize int) {
 		// Rows
 		{
 			query:       `Percentile(field="net_worth", nth=0.5)`,
-			csvVerifier: toCSV("4"),
+			csvVerifier: "10,1\n",
 		},
 	}