feat(firestore): SUM and AVG aggregations

firestore/integration_test.go

@@ -33,6 +33,7 @@ import (
 	apiv1 "cloud.google.com/go/firestore/apiv1/admin"
 	"cloud.google.com/go/firestore/apiv1/admin/adminpb"
 	firestorev1 "cloud.google.com/go/firestore/apiv1/firestorepb"
+	pb "cloud.google.com/go/firestore/apiv1/firestorepb"
 	"cloud.google.com/go/internal/pretty"
 	"cloud.google.com/go/internal/testutil"
 	"cloud.google.com/go/internal/uid"
@@ -43,6 +44,7 @@ import (
 	"google.golang.org/genproto/googleapis/type/latlng"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 func TestMain(m *testing.M) {
@@ -133,7 +135,11 @@ func initIntegrationTest() {
 func createIndexes(ctx context.Context, dbPath string) {
 	var createIndexWg sync.WaitGroup
 
-	indexFields := [][]string{{"updatedAt", "weight", "height"}, {"weight", "height"}}
+	indexFields := [][]string{
+		{"updatedAt", "weight", "height"},
+		{"weight", "height"},
+		{"width", "depth"},
+		{"width", "model"}}
 	indexNames = make([]string, len(indexFields))
 	indexParent := fmt.Sprintf("%s/collectionGroups/%s", dbPath, iColl.ID)
 
@@ -2225,6 +2231,134 @@ func TestIntegration_BulkWriter(t *testing.T) {
 	})
 }
 
+func TestIntegration_AggregationQueries(t *testing.T) {
+	ctx := context.Background()
+	coll := integrationColl(t)
+	h := testHelper{t}
+	docs := []map[string]interface{}{
+		{"width": 1.5, "depth": 99, "model": "A"},
+		{"width": 2.6, "depth": 98, "model": "A"},
+		{"width": 3.7, "depth": 97, "model": "B"},
+		{"width": 4.8, "depth": 96, "model": "B"},
+		{"width": 5.9, "depth": 95, "model": "C"},
+		{"width": 6.0, "depth": 94, "model": "B"},
+		{"width": 7.1, "depth": 93, "model": "C"},
+		{"width": 8.2, "depth": 93, "model": "A"},
+	}
+	for _, doc := range docs {
+		newDoc := coll.NewDoc()
+		h.mustCreate(newDoc, doc)
+	}
+
+	query := coll.Where("width", ">=", 1)
+
+	limitQuery := coll.Where("width", ">=", 1).Limit(4)
+	limitQueryPtr := &limitQuery
+
+	emptyResultsQuery := coll.Where("width", "<", 1)
+	emptyResultsQueryPtr := &emptyResultsQuery
+
+	testcases := []struct {
+		desc             string
+		aggregationQuery *AggregationQuery
+		wantErr          bool
+		result           AggregationResult
+	}{
+		{
+			desc:             "Multiple aggregations",
+			aggregationQuery: query.NewAggregationQuery().WithCount("count1").WithAvg("width", "width_avg1").WithAvg("depth", "depth_avg1").WithSum("width", "width_sum1").WithSum("depth", "depth_sum1"),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"count1":     &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(8)}},
+				"width_sum1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(39.8)}},
+				"depth_sum1": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(765)}},
+				"width_avg1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(4.975)}},
+				"depth_avg1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(95.625)}},
+			},
+		},
+		{
+			desc:             "WithSum aggregation without alias",
+			aggregationQuery: query.NewAggregationQuery().WithSum("width", ""),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"field_1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(39.8)}},
+			},
+		},
+		{
+			desc:             "WithSumPath aggregation without alias",
+			aggregationQuery: query.NewAggregationQuery().WithSumPath([]string{"width"}, ""),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"field_1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(39.8)}},
+			},
+		},
+		{
+			desc:             "WithAvg aggregation without alias",
+			aggregationQuery: query.NewAggregationQuery().WithAvg("width", ""),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"field_1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(4.975)}},
+			},
+		},
+		{
+			desc:             "WithAvgPath aggregation without alias",
+			aggregationQuery: query.NewAggregationQuery().WithAvgPath([]string{"width"}, ""),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"field_1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(4.975)}},
+			},
+		},
+		{
+			desc:             "Aggregations with limit",
+			aggregationQuery: limitQueryPtr.NewAggregationQuery().WithCount("count1").WithAvgPath([]string{"width"}, "width_avg1").WithSumPath([]string{"width"}, "width_sum1"),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"count1":     &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(4)}},
+				"width_sum1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(12.6)}},
+				"width_avg1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(3.15)}},
+			},
+		},
+		{
+			desc:             "Aggregations on empty results",
+			aggregationQuery: emptyResultsQueryPtr.NewAggregationQuery().WithCount("count1").WithAvg("width", "width_avg1").WithSum("width", "width_sum1"),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"count1":     &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(0)}},
+				"width_sum1": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(0)}},
+				"width_avg1": &pb.Value{ValueType: &pb.Value_NullValue{NullValue: structpb.NullValue_NULL_VALUE}},
+			},
+		},
+		{
+			desc:             "Aggregation on non-numeric field",
+			aggregationQuery: query.NewAggregationQuery().WithAvg("model", "model_avg1").WithSum("model", "model_sum1"),
+			wantErr:          false,
+			result: map[string]interface{}{
+				"model_sum1": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(0)}},
+				"model_avg1": &pb.Value{ValueType: &pb.Value_NullValue{NullValue: structpb.NullValue_NULL_VALUE}},
+			},
+		},
+		{
+			desc:             "Aggregation on non existent key",
+			aggregationQuery: query.NewAggregationQuery().WithAvg("randKey", "key_avg1").WithSum("randKey", "key_sum1"),
+			wantErr:          true,
+		},
+	}
+
+	for _, tc := range testcases {
+		aggResult, err := tc.aggregationQuery.Get(ctx)
+		if err != nil && !tc.wantErr {
+			t.Errorf("%s: got: %v, want: nil", tc.desc, err)
+			continue
+		} else if err == nil && tc.wantErr {
+			t.Errorf("%s: got: %v, wanted error", tc.desc, err)
+			continue
+		} else if !reflect.DeepEqual(aggResult, tc.result) {
+			t.Errorf("%s: got: %v, want: %v", tc.desc, aggResult, tc.result)
+			continue
+		}
+	}
+}
+
 func TestIntegration_CountAggregationQuery(t *testing.T) {
 	str := uid.NewSpace("firestore-count", &uid.Options{})
 	datum := str.New()

firestore/query.go

@@ -1233,6 +1233,88 @@ func (a *AggregationQuery) WithCount(alias string) *AggregationQuery {
 	return a
 }
 
+// WithSumPath specifies that the aggregation query should provide a sum of the values
+// of the provided field in the results returned by the underlying Query.
+// The path argument can be a single field or a dot-separated sequence of
+// fields, and must not contain any of the runes "˜*/[]".
+// The alias argument can be empty or a valid Firestore document field name. It can be used
+// as key in the AggregationResult to get the sum value. If alias is empty, Firestore
+// will autogenerate a key.
+func (a *AggregationQuery) WithSumPath(fp FieldPath, alias string) *AggregationQuery {
+	ref, err := fref(fp)
+	if err != nil {
+		a.query.err = err
+		return a
+	}
+
+	aq := &pb.StructuredAggregationQuery_Aggregation{
+		Alias: alias,
+		Operator: &pb.StructuredAggregationQuery_Aggregation_Sum_{
+			Sum: &pb.StructuredAggregationQuery_Aggregation_Sum{
+				Field: ref,
+			},
+		},
+	}
+
+	a.aggregateQueries = append(a.aggregateQueries, aq)
+	return a
+}
+
+// WithSum specifies that the aggregation query should provide a sum of the values
+// of the provided field in the results returned by the underlying Query.
+// The alias argument can be empty or a valid Firestore document field name. It can be used
+// as key in the AggregationResult to get the sum value. If alias is empty, Firestore
+// will autogenerate a key.
+func (a *AggregationQuery) WithSum(path string, alias string) *AggregationQuery {
+	fp, err := parseDotSeparatedString(path)
+	if err != nil {
+		a.query.err = err
+		return a
+	}
+	return a.WithSumPath(fp, alias)
+}
+
+// WithAvgPath specifies that the aggregation query should provide an average of the values
+// of the provided field in the results returned by the underlying Query.
+// The path argument can be a single field or a dot-separated sequence of
+// fields, and must not contain any of the runes "˜*/[]".
+// The alias argument can be empty or a valid Firestore document field name. It can be used
+// as key in the AggregationResult to get the average value. If alias is empty, Firestore
+// will autogenerate a key.
+func (a *AggregationQuery) WithAvgPath(fp FieldPath, alias string) *AggregationQuery {
+	ref, err := fref(fp)
+	if err != nil {
+		a.query.err = err
+		return a
+	}
+
+	aq := &pb.StructuredAggregationQuery_Aggregation{
+		Alias: alias,
+		Operator: &pb.StructuredAggregationQuery_Aggregation_Avg_{
+			Avg: &pb.StructuredAggregationQuery_Aggregation_Avg{
+				Field: ref,
+			},
+		},
+	}
+
+	a.aggregateQueries = append(a.aggregateQueries, aq)
+	return a
+}
+
+// WithAvg specifies that the aggregation query should provide an average of the values
+// of the provided field in the results returned by the underlying Query.
+// The alias argument can be empty or a valid Firestore document field name. It can be used
+// as key in the AggregationResult to get the average value. If alias is empty, Firestore
+// will autogenerate a key.
+func (a *AggregationQuery) WithAvg(path string, alias string) *AggregationQuery {
+	fp, err := parseDotSeparatedString(path)
+	if err != nil {
+		a.query.err = err
+		return a
+	}
+	return a.WithAvgPath(fp, alias)
+}
+
 // Get retrieves the aggregation query results from the service.
 func (a *AggregationQuery) Get(ctx context.Context) (AggregationResult, error) {