From 7a16ffff68e095c142597ad5647e280a8944594a Mon Sep 17 00:00:00 2001
From: Jaden Weiss <jaden@jadendw.dev>
Date: Wed, 8 Jul 2020 12:06:08 -0400
Subject: [PATCH] add UnionRows query

---
 executor.go | 111 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 pql/ast.go  |   5 ++-
 2 files changed, 114 insertions(+), 2 deletions(-)

diff --git a/executor.go b/executor.go
index 2c736badd..42abe099c 100644
--- a/executor.go
+++ b/executor.go
@@ -495,6 +495,111 @@ func (e *executor) execute(ctx context.Context, tx Tx, index string, q *pql.Quer
 	return results, nil
 }
 
+// preprocessQuery expands any calls that need preprocessing.
+// So far, this only needs to process UnionRows.
+func (e *executor) preprocessQuery(ctx context.Context, tx Tx, index string, c *pql.Call, shards []uint64, opt *execOptions) (*pql.Call, error) {
+	switch c.Name {
+	case "UnionRows":
+		// Turn UnionRows(Rows(...)) into Union(Row(...), ...).
+		var rows []*pql.Call
+		for _, child := range c.Children {
+			// Check that we can use the call.
+			switch child.Name {
+			case "Rows":
+			case "TopN":
+			default:
+				return nil, errors.Errorf("cannot use %v as a rows query", child)
+			}
+
+			// Execute the call.
+			rowsResult, err := e.executeCall(ctx, tx, index, child, shards, opt)
+			if err != nil {
+				return nil, err
+			}
+
+			// Turn the results into rows calls.
+			var resultRows []*pql.Call
+			switch rowsResult := rowsResult.(type) {
+			case *PairsField:
+				// Translate pairs into rows calls.
+				for _, p := range rowsResult.Pairs {
+					var val interface{}
+					switch {
+					case p.Key != "":
+						val = p.Key
+					default:
+						val = p.ID
+					}
+					resultRows = append(resultRows, &pql.Call{
+						Name: "Row",
+						Args: map[string]interface{}{
+							rowsResult.Field: val,
+						},
+					})
+				}
+			case RowIDs:
+				// Translate Row IDs into Row calls.
+				for _, id := range rowsResult {
+					resultRows = append(resultRows, &pql.Call{
+						Name: "Row",
+						Args: map[string]interface{}{
+							child.Args["_field"].(string): id,
+						},
+					})
+				}
+			default:
+				return nil, errors.Errorf("unexpected Rows type %T", rowsResult)
+			}
+
+			// Propogate any special properties of the call.
+			switch child.Name {
+			case "Rows":
+				// Propogate "from" time, if set.
+				if v, ok := child.Args["from"]; ok {
+					for _, rowCall := range resultRows {
+						rowCall.Args["from"] = v
+					}
+				}
+
+				// Propogate "to" time, if set.
+				if v, ok := child.Args["to"]; ok {
+					for _, rowCall := range resultRows {
+						rowCall.Args["to"] = v
+					}
+				}
+			}
+
+			rows = append(rows, resultRows...)
+		}
+
+		// Generate a Union call over the rows.
+		return &pql.Call{
+			Name:     "Union",
+			Children: rows,
+		}, nil
+
+	default:
+		// Recurse through child calls.
+		out := make([]*pql.Call, len(c.Children))
+		var changed bool
+		for i, child := range c.Children {
+			res, err := e.preprocessQuery(ctx, tx, index, child, shards, opt)
+			if err != nil {
+				return nil, err
+			}
+			if res != child {
+				changed = true
+			}
+			out[i] = res
+		}
+		if changed {
+			c = c.Clone()
+			c.Children = out
+		}
+		return c, nil
+	}
+}
+
 // executeCall executes a call.
 func (e *executor) executeCall(ctx context.Context, tx Tx, index string, c *pql.Call, shards []uint64, opt *execOptions) (interface{}, error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCall")
@@ -534,6 +639,12 @@ func (e *executor) executeCall(ctx context.Context, tx Tx, index string, c *pql.
 		}
 	}
 
+	// Preprocess the query.
+	c, err := e.preprocessQuery(ctx, tx, index, c, shards, opt)
+	if err != nil {
+		return nil, err
+	}
+
 	// Special handling for mutation and top-n calls.
 	if op, ok := e.additionalCountOps[c.Name]; ok {
 		statFn()
diff --git a/pql/ast.go b/pql/ast.go
index 4195e66bf..9d65bb3c5 100644
--- a/pql/ast.go
+++ b/pql/ast.go
@@ -392,8 +392,9 @@ var callInfoByFunc = map[string]callInfo{
 			"n": int64(0),
 		},
 	},
-	"Union": {allowUnknown: false},
-	"Xor":   {allowUnknown: false},
+	"Union":     {allowUnknown: false},
+	"UnionRows": {allowUnknown: false},
+	"Xor":       {allowUnknown: false},
 
 	// things that take _field
 	"TopN": allowUnderField,