feat: Sum and Avg aggregation feature

google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java

@@ -21,6 +21,7 @@
 
 import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.cloud.datastore.aggregation.AggregationBuilder;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
 
@@ -143,6 +144,18 @@ public Builder addAggregation(Aggregation aggregation) {
       return this;
     }
 
+    public Builder addAggregations(AggregationBuilder<?>... aggregationBuilders) {
+      for (AggregationBuilder<?> builder : aggregationBuilders) {
+        this.aggregations.add(builder.build());
+      }
+      return this;
+    }
+
+    public Builder addAggregations(Aggregation... aggregations) {
+      this.aggregations.addAll(Arrays.asList(aggregations));
+      return this;
+    }
+
     public Builder over(StructuredQuery<?> nestedQuery) {
       this.nestedStructuredQuery = nestedQuery;
       this.mode = STRUCTURED;

google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java

@@ -15,6 +15,9 @@
  */
 package com.google.cloud.datastore;
 
+import static com.google.cloud.datastore.ValueType.DOUBLE;
+import static com.google.cloud.datastore.ValueType.LONG;
+
 import com.google.common.base.MoreObjects;
 import com.google.common.base.MoreObjects.ToStringHelper;
 import java.util.Map;
@@ -24,21 +27,62 @@
 /** Represents a result of an {@link AggregationQuery} query submission. */
 public class AggregationResult {
 
-  private final Map<String, LongValue> properties;
+  private final Map<String, ? extends Value<?>> properties;
 
-  public AggregationResult(Map<String, LongValue> properties) {
+  public AggregationResult(Map<String, ? extends Value<?>> properties) {
     this.properties = properties;
   }
 
   /**
-   * Returns a result value for the given alias.
+   * Returns a result value for the given alias. {@link #getLong(String)} is preferred over this
+   * method, Use {@link #getLong(String)} wherever possible.
    *
    * @param alias A custom alias provided in the query or an autogenerated alias in the form of
    *     'property_\d'
    * @return An aggregation result value for the given alias.
    */
   public Long get(String alias) {
-    return properties.get(alias).get();
+    return getLong(alias);
+  }
+
+  /**
+   * Returns a result value for the given alias.
+   *
+   * @param alias A custom alias provided in the query or an autogenerated alias in the form of
+   *     'property_\d'
+   * @return An aggregation result value for the given alias.
+   */
+  public Long getLong(String alias) {
+    Value<?> value = properties.get(alias);
+    switch (value.getType()) {
+      case DOUBLE:
+        return ((Double) value.get()).longValue();
+      case LONG:
+        return (Long) value.get();
+      default:
+        throw new RuntimeException(
+            String.format("Unsupported type %s received for alias '%s'.", value.getType(), alias));
+    }
+  }
+
+  /**
+   * Returns a result value for the given alias.
+   *
+   * @param alias A custom alias provided in the query or an autogenerated alias in the form of
+   *     'property_\d'
+   * @return An aggregation result value for the given alias.
+   */
+  public Double getDouble(String alias) {
+    Value<?> value = properties.get(alias);
+    switch (value.getType()) {
+      case LONG:
+        return ((Long) value.get()).doubleValue();
+      case DOUBLE:
+        return (Double) value.get();
+      default:
+        throw new RuntimeException(
+            String.format("Unsupported type %s received for alias '%s'.", value.getType(), alias));
+    }
   }
 
   @Override
@@ -61,7 +105,7 @@ public int hashCode() {
   @Override
   public String toString() {
     ToStringHelper toStringHelper = MoreObjects.toStringHelper(this);
-    for (Entry<String, LongValue> entry : properties.entrySet()) {
+    for (Entry<String, ? extends Value<?>> entry : properties.entrySet()) {
       toStringHelper.add(entry.getKey(), entry.getValue().get());
     }
     return toStringHelper.toString();

google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java

@@ -38,8 +38,28 @@ public String getAlias() {
   @InternalApi
   public abstract AggregationQuery.Aggregation toPb();
 
+  @InternalApi
+  protected AggregationQuery.Aggregation.Builder aggregationBuilder() {
+    AggregationQuery.Aggregation.Builder aggregationBuilder =
+        AggregationQuery.Aggregation.newBuilder();
+    if (this.getAlias() != null) {
+      aggregationBuilder.setAlias(this.getAlias());
+    }
+    return aggregationBuilder;
+  }
+
   /** Returns a {@link CountAggregation} builder. */
   public static CountAggregation.Builder count() {
     return new CountAggregation.Builder();
   }
+
+  /** Returns a {@link SumAggregation} builder. */
+  public static SumAggregation.Builder sum(String propertyReference) {
+    return new SumAggregation.Builder().propertyReference(propertyReference);
+  }
+
+  /** Returns a {@link AvgAggregation} builder. */
+  public static AvgAggregation.Builder avg(String propertyReference) {
+    return new AvgAggregation.Builder().propertyReference(propertyReference);
+  }
 }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AvgAggregation.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.aggregation;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import com.google.api.core.InternalApi;
+import com.google.datastore.v1.AggregationQuery;
+import com.google.datastore.v1.AggregationQuery.Aggregation.Avg;
+import com.google.datastore.v1.PropertyReference;
+import java.util.Objects;
+
+/** Represents an {@link Aggregation} which returns average of numerical values. */
+public class AvgAggregation extends Aggregation {
+
+  private final String propertyReference;
+
+  public AvgAggregation(String alias, String propertyReference) {
+    super(alias);
+    checkArgument(propertyReference != null, "Property reference can't be null");
+    this.propertyReference = propertyReference;
+  }
+
+  @InternalApi
+  @Override
+  public AggregationQuery.Aggregation toPb() {
+    PropertyReference reference =
+        PropertyReference.newBuilder().setName(this.propertyReference).build();
+    Avg avg = Avg.newBuilder().setProperty(reference).build();
+    return aggregationBuilder().setAvg(avg).build();
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    AvgAggregation that = (AvgAggregation) o;
+    return Objects.equals(this.propertyReference, that.propertyReference)
+        && Objects.equals(getAlias(), that.getAlias());
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(getAlias(), this.propertyReference);
+  }
+
+  /** A builder class to create and customize a {@link AvgAggregation}. */
+  public static class Builder implements AggregationBuilder<AvgAggregation> {
+
+    private String alias;
+    private String propertyReference;
+
+    public AvgAggregation.Builder propertyReference(String propertyReference) {
+      this.propertyReference = propertyReference;
+      return this;
+    }
+
+    public AvgAggregation.Builder as(String alias) {
+      this.alias = alias;
+      return this;
+    }
+
+    @Override
+    public AvgAggregation build() {
+      return new AvgAggregation(alias, propertyReference);
+    }
+  }
+}

google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java

@@ -30,14 +30,7 @@ public CountAggregation(String alias) {
 
   @Override
   public AggregationQuery.Aggregation toPb() {
-    Count.Builder countBuilder = Count.newBuilder();
-
-    AggregationQuery.Aggregation.Builder aggregationBuilder =
-        AggregationQuery.Aggregation.newBuilder().setCount(countBuilder);
-    if (this.getAlias() != null) {
-      aggregationBuilder.setAlias(this.getAlias());
-    }
-    return aggregationBuilder.build();
+    return aggregationBuilder().setCount(Count.newBuilder()).build();
   }
 
   @Override
@@ -49,13 +42,7 @@ public boolean equals(Object o) {
       return false;
     }
     CountAggregation that = (CountAggregation) o;
-    boolean bothAliasAreNull = getAlias() == null && that.getAlias() == null;
-    if (bothAliasAreNull) {
-      return true;
-    } else {
-      boolean bothArePresent = getAlias() != null && that.getAlias() != null;
-      return bothArePresent && getAlias().equals(that.getAlias());
-    }
+    return Objects.equals(getAlias(), that.getAlias());
   }
 
   @Override

google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/SumAggregation.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.aggregation;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import com.google.api.core.InternalApi;
+import com.google.datastore.v1.AggregationQuery;
+import com.google.datastore.v1.AggregationQuery.Aggregation.Sum;
+import com.google.datastore.v1.PropertyReference;
+import java.util.Objects;
+
+/** Represents an {@link Aggregation} which returns sum of numerical values. */
+public class SumAggregation extends Aggregation {
+
+  private final String propertyReference;
+
+  public SumAggregation(String alias, String propertyReference) {
+    super(alias);
+    checkArgument(propertyReference != null, "Property reference can't be null");
+    this.propertyReference = propertyReference;
+  }
+
+  @InternalApi
+  @Override
+  public AggregationQuery.Aggregation toPb() {
+    PropertyReference reference =
+        PropertyReference.newBuilder().setName(this.propertyReference).build();
+    Sum sum = Sum.newBuilder().setProperty(reference).build();
+    return aggregationBuilder().setSum(sum).build();
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    SumAggregation that = (SumAggregation) o;
+    return Objects.equals(this.propertyReference, that.propertyReference)
+        && Objects.equals(getAlias(), that.getAlias());
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(getAlias(), this.propertyReference);
+  }
+
+  /** A builder class to create and customize a {@link SumAggregation}. */
+  public static class Builder implements AggregationBuilder<SumAggregation> {
+
+    private String alias;
+    private String propertyReference;
+
+    public SumAggregation.Builder propertyReference(String propertyReference) {
+      this.propertyReference = propertyReference;
+      return this;
+    }
+
+    public SumAggregation.Builder as(String alias) {
+      this.alias = alias;
+      return this;
+    }
+
+    @Override
+    public SumAggregation build() {
+      return new SumAggregation(alias, propertyReference);
+    }
+  }
+}

google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java

@@ -19,7 +19,6 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.AggregationResults;
-import com.google.cloud.datastore.LongValue;
 import com.google.datastore.v1.RunAggregationQueryResponse;
 import com.google.datastore.v1.Value;
 import java.util.AbstractMap.SimpleEntry;
@@ -39,20 +38,19 @@ public AggregationResults transform(RunAggregationQueryResponse response) {
     Timestamp readTime = Timestamp.fromProto(response.getBatch().getReadTime());
     List<AggregationResult> aggregationResults =
         response.getBatch().getAggregationResultsList().stream()
-            .map(
-                aggregationResult -> new AggregationResult(resultWithLongValues(aggregationResult)))
+            .map(aggregationResult -> new AggregationResult(transformValues(aggregationResult)))
             .collect(Collectors.toCollection(LinkedList::new));
     return new AggregationResults(aggregationResults, readTime);
   }
 
-  private Map<String, LongValue> resultWithLongValues(
+  private Map<String, com.google.cloud.datastore.Value<?>> transformValues(
       com.google.datastore.v1.AggregationResult aggregationResult) {
     return aggregationResult.getAggregatePropertiesMap().entrySet().stream()
         .map(
-            (Function<Entry<String, Value>, Entry<String, LongValue>>)
+            (Function<Entry<String, Value>, Entry<String, com.google.cloud.datastore.Value<?>>>)
                 entry ->
                     new SimpleEntry<>(
-                        entry.getKey(), (LongValue) LongValue.fromPb(entry.getValue())))
+                        entry.getKey(), com.google.cloud.datastore.Value.fromPb(entry.getValue())))
         .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
   }
 }