feat(datastore): SUM and AVG aggregations

datastore/integration_test.go

@@ -37,6 +37,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 // TODO(djd): Make test entity clean up more robust: some test entities may
@@ -471,6 +472,8 @@ func TestIntegration_NilKey(t *testing.T) {
 type SQChild struct {
 	I, J int
 	T, U int64
+	V    float64
+	W    string
 }
 
 type SQTestCase struct {
@@ -701,17 +704,17 @@ func TestIntegration_AggregationQueries(t *testing.T) {
 	client := newTestClient(ctx, t)
 	defer client.Close()
 
-	parent := NameKey("SQParent", "TestIntegration_Filters"+suffix, nil)
+	parent := NameKey("SQParent", "TestIntegration_AggregationQueries"+suffix, nil)
 	now := timeNow.Truncate(time.Millisecond).Unix()
 	children := []*SQChild{
-		{I: 0, T: now, U: now},
-		{I: 1, T: now, U: now},
-		{I: 2, T: now, U: now},
-		{I: 3, T: now, U: now},
-		{I: 4, T: now, U: now},
-		{I: 5, T: now, U: now},
-		{I: 6, T: now, U: now},
-		{I: 7, T: now, U: now},
+		{I: 0, T: now, U: now, V: 1.5, W: "str"},
+		{I: 1, T: now, U: now, V: 1.5, W: "str"},
+		{I: 2, T: now, U: now, V: 1.5, W: "str"},
+		{I: 3, T: now, U: now, V: 1.5, W: "str"},
+		{I: 4, T: now, U: now, V: 1.5, W: "str"},
+		{I: 5, T: now, U: now, V: 1.5, W: "str"},
+		{I: 6, T: now, U: now, V: 1.5, W: "str"},
+		{I: 7, T: now, U: now, V: 1.5, W: "str"},
 	}
 
 	keys := make([]*Key, len(children))
@@ -739,20 +742,75 @@ func TestIntegration_AggregationQueries(t *testing.T) {
 	}{
 		{
 			desc:          "Count Failure - Missing index",
-			aggQuery:      baseQuery.Filter("T>=", now).NewAggregationQuery().WithCount("count"),
+			aggQuery:      NewQuery("SQChild").Ancestor(parent).Filter("T>=", now).NewAggregationQuery().WithCount("count"),
 			wantFailure:   true,
 			wantErrMsg:    "no matching index found",
 			wantAggResult: nil,
 		},
 		{
 			desc:        "Count Success",
-			aggQuery:    baseQuery.Filter("T=", now).Filter("I>=", 3).NewAggregationQuery().WithCount("count"),
+			aggQuery:    NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Filter("I>=", 3).NewAggregationQuery().WithCount("count"),
 			wantFailure: false,
 			wantErrMsg:  "",
 			wantAggResult: map[string]interface{}{
 				"count": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 5}},
 			},
 		},
+		{
+			desc:        "Multiple aggregations",
+			aggQuery:    NewQuery("SQChild").Ancestor(parent).Filter("T=", now).NewAggregationQuery().WithSum("I", "i_sum").WithAvg("I", "avg").WithSum("V", "v_sum"),
+			wantFailure: false,
+			wantErrMsg:  "",
+			wantAggResult: map[string]interface{}{
+				"i_sum": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 28}},
+				"v_sum": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: 12}},
+				"avg":   &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: 3.5}},
+			},
+		},
+		{
+			desc:        "Multiple aggregations with limit ",
+			aggQuery:    NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Limit(2).NewAggregationQuery().WithSum("I", "sum").WithAvg("I", "avg"),
+			wantFailure: false,
+			wantErrMsg:  "",
+			wantAggResult: map[string]interface{}{
+				"sum": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 1}},
+				"avg": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: 0.5}},
+			},
+		},
+		{
+			desc:        "Multiple aggregations on non-numeric field",
+			aggQuery:    NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Limit(2).NewAggregationQuery().WithSum("W", "sum").WithAvg("W", "avg"),
+			wantFailure: false,
+			wantErrMsg:  "",
+			wantAggResult: map[string]interface{}{
+				"sum": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(0)}},
+				"avg": &pb.Value{ValueType: &pb.Value_NullValue{NullValue: structpb.NullValue_NULL_VALUE}},
+			},
+		},
+		{
+			desc:        "Sum aggregation without alias",
+			aggQuery:    NewQuery("SQChild").Ancestor(parent).Filter("T=", now).NewAggregationQuery().WithSum("I", ""),
+			wantFailure: false,
+			wantErrMsg:  "",
+			wantAggResult: map[string]interface{}{
+				"property_1": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 28}},
+			},
+		},
+		{
+			desc:        "Average aggregation without alias",
+			aggQuery:    NewQuery("SQChild").Ancestor(parent).Filter("T=", now).NewAggregationQuery().WithAvg("I", ""),
+			wantFailure: false,
+			wantErrMsg:  "",
+			wantAggResult: map[string]interface{}{
+				"property_1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: 3.5}},
+			},
+		},
+		{
+			desc:        "Aggregation with invalid field name",
+			aggQuery:    NewQuery("SQChild").Ancestor(parent).Filter("T=", now).NewAggregationQuery().WithSum("I2", "").WithAvg("I2", ""),
+			wantFailure: true,
+			wantErrMsg:  "",
+		},
 	}
 
 	for _, testCase := range testCases {

datastore/query.go

@@ -1057,5 +1057,49 @@ func (aq *AggregationQuery) WithCount(alias string) *AggregationQuery {
 	return aq
 }
 
+// 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 Datastore entity property name. It can be used
+// as key in the AggregationResult to get the sum value. If alias is empty, Datastore
+// will autogenerate a key.
+func (aq *AggregationQuery) WithSum(fieldName string, alias string) *AggregationQuery {
+	aqpb := &pb.AggregationQuery_Aggregation{
+		Alias: alias,
+		Operator: &pb.AggregationQuery_Aggregation_Sum_{
+			Sum: &pb.AggregationQuery_Aggregation_Sum{
+				Property: &pb.PropertyReference{
+					Name: fieldName,
+				},
+			},
+		},
+	}
+
+	aq.aggregationQueries = append(aq.aggregationQueries, aqpb)
+
+	return aq
+}
+
+// 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 Datastore entity property name. It can be used
+// as key in the AggregationResult to get the sum value. If alias is empty, Datastore
+// will autogenerate a key.
+func (aq *AggregationQuery) WithAvg(fieldName string, alias string) *AggregationQuery {
+	aqpb := &pb.AggregationQuery_Aggregation{
+		Alias: alias,
+		Operator: &pb.AggregationQuery_Aggregation_Avg_{
+			Avg: &pb.AggregationQuery_Aggregation_Avg{
+				Property: &pb.PropertyReference{
+					Name: fieldName,
+				},
+			},
+		},
+	}
+
+	aq.aggregationQueries = append(aq.aggregationQueries, aqpb)
+
+	return aq
+}
+
 // AggregationResult contains the results of an aggregation query.
 type AggregationResult map[string]interface{}

datastore/testdata/index.yaml

@@ -38,4 +38,17 @@ indexes:
   properties:
   - name: T
   - name: J
-  - name: U
+  - name: U
+
+- kind: SQChild
+  ancestor: yes
+  properties:
+  - name: T
+  - name: I
+  - name: V
+
+- kind: SQChild
+  ancestor: yes
+  properties:
+  - name: T
+  - name: W