Merge pull request #1961 from molecula/fb-1227

[FB-1227] Enable WHERE clause for SELECT DISTINCT SQL queries

sql/query.go

@@ -98,8 +98,14 @@ func Between(fieldName string, a interface{}, b interface{}) string {
 }
 
 // Distinct creates a Distinct query.
-func Distinct(indexName, fieldName string) string {
-	return fmt.Sprintf("Distinct(Row(%s!=null),index='%s',field='%s')", fieldName, indexName, fieldName)
+func Distinct(indexName, fieldName, rowCall string) string {
+	var b strings.Builder
+	fmt.Fprintf(&b, `Distinct(`)
+	if rowCall != "" {
+		fmt.Fprintf(&b, `%s, `, rowCall)
+	}
+	fmt.Fprintf(&b, `index='%s',field='%s')`, indexName, fieldName)
+	return b.String()
 }
 
 // RowDistinct creates a Distinct query with the given row filter.

sql/reduce.go

@@ -347,6 +347,34 @@ func AssignHeaders(rowser pproto.ToRowser, headers ...Column) pproto.ToRowser {
 	return &assignHeadersRowser{rowser, headers}
 }
 
+type staticHeaderRowser struct {
+	rowser pproto.ToRowser
+	cols   []Column
+}
+
+func (a *staticHeaderRowser) ToRows(fn func(*pproto.RowResponse) error) error {
+	return a.rowser.ToRows(func(row *pproto.RowResponse) error {
+		var out pproto.RowResponse
+
+		headers := make([]*pproto.ColumnInfo, len(row.Headers))
+		for i := range row.Headers {
+			header := row.Headers[i]
+			header.Name = a.cols[i].Name()
+			headers[i] = header
+		}
+		out.Headers = headers
+
+		out.Columns = row.Columns
+
+		return fn(&out)
+	})
+}
+
+// StaticHeaders assigns fixed cols to a ToRowser.
+func StaticHeaders(rowser pproto.ToRowser, cols ...Column) pproto.ToRowser {
+	return &staticHeaderRowser{rowser, cols}
+}
+
 var (
 	ErrIncompleteHeaders = errors.New("incomplete header assignment")
 	ErrFieldNotInHeaders = errors.New("field not found in source header")

sql/router.go

@@ -29,6 +29,17 @@ func newRouter() *router {
 		handlerSelectFieldsFromTableWhere{},
 	)
 	////
+	selectRouter.addFilter(
+		NewQueryMask(
+			SelectPartDistinct|SelectPartField,
+			FromPartTable,
+			WherePartFieldCondition|WherePartMultiFieldCondition,
+			0,
+			0,
+		),
+		[]QueryMask{},
+		handlerSelectDistinctFromTable{},
+	)
 	selectRouter.addRoute("select distinct fld from tbl", handlerSelectDistinctFromTable{})
 	////
 	selectRouter.addFilter(

sql/select.go

@@ -305,6 +305,15 @@ func (h handlerSelectDistinctFromTable) Apply(stmt *sqlparser.Select, qm QueryMa
 		return nil, errors.New("distinct requires a valid field column")
 	}
 
+	var wherePQL string
+	if stmt.Where != nil {
+		if wherePQL, err = extractWhere(index, stmt.Where.Expr); err != nil {
+			return nil, err
+		}
+	} else {
+		wherePQL = All()
+	}
+
 	limit, offset, hasLimit, hasOffset, err := extractLimitOffset(stmt)
 	if err != nil {
 		return nil, errors.Wrap(err, "extracting limit")
@@ -315,22 +324,8 @@ func (h handlerSelectDistinctFromTable) Apply(stmt *sqlparser.Select, qm QueryMa
 		return nil, errors.Wrap(err, "extracting order by")
 	}
 
-	// Determine the type of the field needing distinct.
-	// If the pilosa field is type int, handle it as a Distinct() query.
-	// Otherwise, use Rows()
-	// TODO: ensure this works for all field types (bool, time, etc).
-	var qo string
-	if fieldCol.Field.Type() == pilosa.FieldTypeInt || fieldCol.Field.Type() == pilosa.FieldTypeTimestamp {
-		qo = Distinct(fieldCol.Field.Index(), fieldCol.Field.Name())
-	} else {
-		if !qm.HasOrderBy() && limit > 0 {
-			if qo, err = RowsLimit(fieldCol.Field.Name(), int64(limit)); err != nil {
-				return nil, errors.Wrap(err, "creating Rows query")
-			}
-		} else {
-			qo = Rows(fieldCol.Field.Name())
-		}
-	}
+	// We use a Distinct call instead of Rows as it supports filtering.
+	qo := Distinct(fieldCol.Field.Index(), fieldCol.Field.Name(), wherePQL)
 
 	mr := &MappingResult{
 		IndexName: indexName,
@@ -340,7 +335,7 @@ func (h handlerSelectDistinctFromTable) Apply(stmt *sqlparser.Select, qm QueryMa
 
 	// Assign headers to the result.
 	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
-		return AssignHeaders(result, selectFields...)
+		return StaticHeaders(result, selectFields...)
 	})
 
 	if qm.HasOrderBy() {
@@ -795,7 +790,7 @@ func (h handlerSelectJoin) Apply(stmt *sqlparser.Select, qm QueryMask, indexFunc
 	// Build the Distinct() portion of the query on the secondary.
 	var distinctQry string
 	if secondaryWhere == "" {
-		distinctQry = Distinct(secondaryField.Index(), secondaryField.Name())
+		distinctQry = Distinct(secondaryField.Index(), secondaryField.Name(), "")
 	} else {
 		distinctQry = RowDistinct(secondaryField.Index(), secondaryField.Name(), secondaryWhere)
 	}