OGM-1542 Distinct count implementation

hibernate · Dec 14, 2018 · 896f5f8 · 896f5f8
1 parent 1ece625
commit 896f5f8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 16 deletions.
diff --git a/mongodb/src/main/java/org/hibernate/ogm/datastore/mongodb/MongoDBDialect.java b/mongodb/src/main/java/org/hibernate/ogm/datastore/mongodb/MongoDBDialect.java
@@ -1037,7 +1037,7 @@ private static ClosableIterator<Tuple> doAggregate(MongoDBQueryDescriptor query,
 		pipeline.add( stage( "$match", query.getCriteria() ) );
 
 		if ( query.getAggregation() != null ) {
-			pipeline.add( query.getAggregation().asDocument() );
+			pipeline.addAll( query.getAggregation().asDocumentPipeline() );
 		}
 
 		if ( query.getProjection() != null ) {

diff --git a/...main/java/org/hibernate/ogm/datastore/mongodb/query/parsing/impl/AggregationRenderer.java b/...main/java/org/hibernate/ogm/datastore/mongodb/query/parsing/impl/AggregationRenderer.java
@@ -6,6 +6,9 @@
  */
 package org.hibernate.ogm.datastore.mongodb.query.parsing.impl;
 
+import java.util.Arrays;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.Locale;
 
 import org.bson.Document;
@@ -27,14 +30,14 @@
  *     {$sum: $col }}
  * }}
  *
- * @see #asDocument()
+ * @see #asDocumentPipeline()
  * @author Davide D'Alto
  * @author Aleksandr Mylnikov
  */
 public class AggregationRenderer {
 
 	private static final String PROJECTION_FIELD = "n";
-	private final boolean isCount;
+	private final AggregationPropertyPath.Type aggregationType;
 	private final String propertyPath;
 	private final String aggregationTypeOperator;
 
@@ -43,24 +46,20 @@ public AggregationRenderer(AggregationPropertyPath.Type aggregationType) {
 	}
 
 	public AggregationRenderer(String propertypath, AggregationPropertyPath.Type aggregationType) {
-		validate( aggregationType );
-		this.propertyPath = "$" + propertypath;
-		this.aggregationTypeOperator = "$" + aggregationType.name().toLowerCase( Locale.ROOT );
-		this.isCount = aggregationType == Type.COUNT;
+		this.propertyPath = propertypath != null ? "$" + propertypath : null;
+		this.aggregationTypeOperator = "$" + aggregationType.name().toLowerCase( Locale.ROOT ).split( "_" )[0];
+		this.aggregationType = aggregationType;
 	}
 
-	private void validate(Type aggregationType) {
-		if ( aggregationType == Type.COUNT_DISTINCT ) {
-			throw new UnsupportedOperationException( "Currently OGM does not support count distinct operations" );
+	public List<Document> asDocumentPipeline() {
+		if ( aggregationType == Type.COUNT_DISTINCT && propertyPath != null ) {
+			return getDistinctCount();
 		}
-	}
-
-	public Document asDocument() {
-		if ( isCount ) {
-			return getCount();
+		else if ( aggregationType == Type.COUNT ) {
+			return Arrays.asList( getCount() );
 		}
 		else {
-			return getGroup();
+			return Arrays.asList( getGroup() );
 		}
 	}
 
@@ -73,6 +72,13 @@ private Document getGroup() {
 		return group;
 	}
 
+	private List<Document> getDistinctCount() {
+		List<Document> distinctCount = new LinkedList<>();
+		distinctCount.add( new Document().append( "$group", new Document().append( "_id", propertyPath ) ) );
+		distinctCount.add( getCount() );
+		return distinctCount;
+	}
+
 	private Document getCount() {
 		Document count = new Document();
 		count.append( aggregationTypeOperator, PROJECTION_FIELD );