Skip to content

Commit 6cb9138

Browse files
kotharironakkotharironak
authored
feat: add supports for selection expression and group_by clause for new query api (#93)
* feat: add support for aggregation selection * feat: add support for aggregation selecition and groupby clause * removes additional comments
1 parent 04a5518 commit 6cb9138

File tree

9 files changed

+256
-10
lines changed

9 files changed

+256
-10
lines changed

‎document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,16 @@
11
package org.hypertrace.core.documentstore;
22

3+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.AVG;
4+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.COUNT;
5+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.DISTINCT_COUNT;
6+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.MAX;
7+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.MIN;
8+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.SUM;
39
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.MULTIPLY;
410
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.AND;
511
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.OR;
612
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
13+
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.GT;
714
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.GTE;
815
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.LTE;
916
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NEQ;
@@ -37,6 +44,7 @@
3744
import org.apache.commons.lang3.tuple.ImmutablePair;
3845
import org.bson.codecs.configuration.CodecConfigurationException;
3946
import org.hypertrace.core.documentstore.Filter.Op;
47+
import org.hypertrace.core.documentstore.expression.impl.AggregateExpression;
4048
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
4149
import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
4250
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
@@ -1641,6 +1649,44 @@ public void testQueryV1FunctionalSelectionExpressionWithNestedFieldWithAlias(Str
16411649
"mongo/test_selection_expression_nested_fields_alias_result.json");
16421650
}
16431651

1652+
@ParameterizedTest
1653+
@MethodSource("databaseContextProvider")
1654+
public void testQueryV1AggregationExpression(String dataStoreName) throws IOException {
1655+
Map<Key, Document> documents = createDocumentsFromResource("mongo/collection_data.json");
1656+
Datastore datastore = datastoreMap.get(dataStoreName);
1657+
Collection collection = datastore.getCollection(COLLECTION_NAME);
1658+
1659+
// add docs
1660+
boolean result = collection.bulkUpsert(documents);
1661+
Assertions.assertTrue(result);
1662+
1663+
org.hypertrace.core.documentstore.query.Query query =
1664+
org.hypertrace.core.documentstore.query.Query.builder()
1665+
.setFilter(
1666+
RelationalExpression.of(
1667+
IdentifierExpression.of("price"), GT, ConstantExpression.of(5)))
1668+
.addSelection(IdentifierExpression.of("item"))
1669+
.addSelection(
1670+
AggregateExpression.of(AVG, IdentifierExpression.of("quantity")), "qty_avg")
1671+
.addSelection(
1672+
AggregateExpression.of(COUNT, IdentifierExpression.of("quantity")), "qty_count")
1673+
.addSelection(
1674+
AggregateExpression.of(DISTINCT_COUNT, IdentifierExpression.of("quantity")),
1675+
"qty_distinct_count")
1676+
.addSelection(
1677+
AggregateExpression.of(SUM, IdentifierExpression.of("quantity")), "qty_sum")
1678+
.addSelection(
1679+
AggregateExpression.of(MIN, IdentifierExpression.of("quantity")), "qty_min")
1680+
.addSelection(
1681+
AggregateExpression.of(MAX, IdentifierExpression.of("quantity")), "qty_max")
1682+
.addAggregation(IdentifierExpression.of("item"))
1683+
.build();
1684+
1685+
Iterator<Document> resultDocs = collection.aggregate(query);
1686+
assertSizeAndDocsEqual(
1687+
dataStoreName, resultDocs, 3, "mongo/test_aggregation_expression_result.json");
1688+
}
1689+
16441690
@ParameterizedTest
16451691
@MethodSource("databaseContextProvider")
16461692
public void whenBulkUpdatingNonExistentRecords_thenExpectNothingToBeUpdatedOrCreated(

‎document-store/src/integrationTest/resources/mongo/test_aggregation_expression_result.json‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
[
2+
{
3+
"qty_min":1,
4+
"qty_avg":1.0,
5+
"item":"Mirror",
6+
"qty_max":1,
7+
"qty_count":1,
8+
"qty_distinct_count":1,
9+
"qty_sum":1
10+
},
11+
{
12+
"qty_min":5,
13+
"qty_avg":7.5,
14+
"item":"Comb",
15+
"qty_max":10,
16+
"qty_count":2,
17+
"qty_distinct_count":2,
18+
"qty_sum":15
19+
},
20+
{
21+
"qty_min":2,
22+
"qty_avg":4.0,
23+
"item":"Soap",
24+
"qty_max":5,
25+
"qty_count":3,
26+
"qty_distinct_count":2,
27+
"qty_sum":12
28+
}
29+
]

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -343,9 +343,10 @@ private CloseableIterator<Document> executeQueryV1(
343343
: new PostgresResultIterator(resultSet);
344344
return closeableIterator;
345345
} catch (SQLException e) {
346-
LOGGER.error("SQLException querying documents. query: {}", query, e);
346+
LOGGER.error(
347+
"SQLException querying documents. original query: {}, sql query:", query, sqlQuery, e);
348+
throw new RuntimeException(e);
347349
}
348-
return EMPTY_ITERATOR;
349350
}
350351

351352
@VisibleForTesting

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParser.java‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@
33
import java.util.List;
44
import java.util.Optional;
55
import org.hypertrace.core.documentstore.expression.type.FilterTypeExpression;
6-
import org.hypertrace.core.documentstore.expression.type.GroupTypeExpression;
76
import org.hypertrace.core.documentstore.postgres.Params;
87
import org.hypertrace.core.documentstore.postgres.Params.Builder;
98
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresFilterTypeExpressionVisitor;
9+
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresGroupTypeExpressionVisitor;
1010
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSelectTypeExpressionVisitor;
1111
import org.hypertrace.core.documentstore.query.Pagination;
1212
import org.hypertrace.core.documentstore.query.Query;
@@ -45,7 +45,7 @@ public String parse(Query query) {
4545
}
4646

4747
// group by
48-
Optional<String> groupBy = parseGroupBy(query.getAggregations());
48+
Optional<String> groupBy = parseGroupBy(query);
4949
if (groupBy.isPresent()) {
5050
sqlBuilder.append(String.format(" GROUP BY %s", groupBy.get()));
5151
}
@@ -80,11 +80,8 @@ private Optional<String> parseFilter(Optional<FilterTypeExpression> filterTypeEx
8080
expression -> expression.accept(new PostgresFilterTypeExpressionVisitor(this)));
8181
}
8282

83-
private Optional<String> parseGroupBy(List<GroupTypeExpression> groupTypeExpressionList) {
84-
if (groupTypeExpressionList.size() > 0) {
85-
throw new UnsupportedOperationException(String.format(NOT_YET_SUPPORTED, "group by clause"));
86-
}
87-
return Optional.empty();
83+
private Optional<String> parseGroupBy(Query query) {
84+
return Optional.ofNullable(PostgresGroupTypeExpressionVisitor.getGroupByClause(query));
8885
}
8986

9087
private Optional<String> parseHaving(Optional<FilterTypeExpression> filterTypeExpression) {

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresAggregateExpressionVisitor.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
package org.hypertrace.core.documentstore.postgres.query.v1.vistors;
2+
3+
import lombok.NoArgsConstructor;
4+
import org.hypertrace.core.documentstore.expression.impl.AggregateExpression;
5+
import org.hypertrace.core.documentstore.expression.operators.AggregationOperator;
6+
7+
@NoArgsConstructor
8+
public class PostgresAggregateExpressionVisitor extends PostgresSelectTypeExpressionVisitor {
9+
10+
public PostgresAggregateExpressionVisitor(PostgresSelectTypeExpressionVisitor baseVisitor) {
11+
super(baseVisitor);
12+
}
13+
14+
@Override
15+
public String visit(final AggregateExpression expression) {
16+
AggregationOperator operator = expression.getAggregator();
17+
PostgresSelectTypeExpressionVisitor selectTypeExpressionVisitor =
18+
new PostgresFunctionExpressionVisitor(
19+
new PostgresDataAccessorIdentifierExpressionVisitor(
20+
new PostgresConstantExpressionVisitor(this)));
21+
22+
String value = expression.getExpression().accept(selectTypeExpressionVisitor);
23+
return value != null ? convertToAggregationFunction(operator, value) : null;
24+
}
25+
26+
private String convertToAggregationFunction(AggregationOperator operator, String value) {
27+
if (operator.equals(AggregationOperator.DISTINCT_COUNT)) {
28+
return String.format("COUNT(DISTINCT %s )", value);
29+
}
30+
return String.format("%s( %s )", operator, value);
31+
}
32+
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,18 @@
33
import java.util.Objects;
44
import java.util.stream.Collector;
55
import java.util.stream.Collectors;
6+
import lombok.NoArgsConstructor;
67
import org.apache.commons.lang3.StringUtils;
78
import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
89
import org.hypertrace.core.documentstore.expression.operators.FunctionOperator;
910

11+
@NoArgsConstructor
1012
public class PostgresFunctionExpressionVisitor extends PostgresSelectTypeExpressionVisitor {
13+
14+
public PostgresFunctionExpressionVisitor(PostgresSelectTypeExpressionVisitor baseVisitor) {
15+
super(baseVisitor);
16+
}
17+
1118
@Override
1219
public String visit(final FunctionExpression expression) {
1320
int numArgs = expression.getOperands().size();

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresGroupTypeExpressionVisitor.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
package org.hypertrace.core.documentstore.postgres.query.v1.vistors;
2+
3+
import java.util.List;
4+
import java.util.Objects;
5+
import java.util.stream.Collectors;
6+
import org.apache.commons.lang3.StringUtils;
7+
import org.hypertrace.core.documentstore.expression.impl.AggregateExpression;
8+
import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
9+
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
10+
import org.hypertrace.core.documentstore.expression.operators.AggregationOperator;
11+
import org.hypertrace.core.documentstore.expression.type.GroupTypeExpression;
12+
import org.hypertrace.core.documentstore.parser.GroupTypeExpressionVisitor;
13+
import org.hypertrace.core.documentstore.query.Query;
14+
15+
public class PostgresGroupTypeExpressionVisitor implements GroupTypeExpressionVisitor {
16+
17+
PostgresSelectTypeExpressionVisitor selectTypeExpressionVisitor =
18+
new PostgresFieldIdentifierExpressionVisitor();
19+
20+
@Override
21+
public String visit(final FunctionExpression expression) {
22+
throw new UnsupportedOperationException(
23+
"FunctionalExpression on Group by is not yet supported");
24+
}
25+
26+
@Override
27+
public String visit(final IdentifierExpression expression) {
28+
return selectTypeExpressionVisitor.visit(expression);
29+
}
30+
31+
public static String getGroupByClause(final Query query) {
32+
if (query.getAggregations().size() <= 0) return null;
33+
34+
if (!validate(query)) {
35+
throw new UnsupportedOperationException("Group By clause with DISTINCT is not yet supported");
36+
}
37+
38+
List<GroupTypeExpression> groupTypeExpressions = query.getAggregations();
39+
40+
PostgresGroupTypeExpressionVisitor groupTypeExpressionVisitor =
41+
new PostgresGroupTypeExpressionVisitor();
42+
43+
String childList =
44+
groupTypeExpressions.stream()
45+
.map(exp -> exp.accept(groupTypeExpressionVisitor))
46+
.filter(Objects::nonNull)
47+
.map(Object::toString)
48+
.filter(StringUtils::isNotEmpty)
49+
.collect(Collectors.joining(","));
50+
51+
return !childList.isEmpty() ? childList : null;
52+
}
53+
54+
private static boolean validate(Query query) {
55+
return !query.getSelections().stream()
56+
.filter(exp -> exp.getExpression() instanceof AggregateExpression)
57+
.anyMatch(
58+
exp ->
59+
((AggregateExpression) exp.getExpression())
60+
.getAggregator()
61+
.equals(AggregationOperator.DISTINCT));
62+
}
63+
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresSelectTypeExpressionVisitor.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,8 @@ static class PgSelection {
5353

5454
public static String getSelections(final List<SelectionSpec> selectionSpecs) {
5555
PostgresSelectTypeExpressionVisitor selectTypeExpressionVisitor =
56-
new PostgresFieldIdentifierExpressionVisitor(new PostgresFunctionExpressionVisitor());
56+
new PostgresAggregateExpressionVisitor(
57+
new PostgresFieldIdentifierExpressionVisitor(new PostgresFunctionExpressionVisitor()));
5758

5859
// used for if alias is missing
5960
PostgresIdentifierExpressionVisitor identifierExpressionVisitor =

‎document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,21 @@
11
package org.hypertrace.core.documentstore.postgres.query.v1;
22

3+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.AVG;
4+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.COUNT;
5+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.DISTINCT;
6+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.DISTINCT_COUNT;
7+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.MAX;
8+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.MIN;
9+
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.SUM;
310
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.MULTIPLY;
411
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.AND;
512
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.OR;
13+
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
614
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.GTE;
715
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.LTE;
816
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NEQ;
917

18+
import org.hypertrace.core.documentstore.expression.impl.AggregateExpression;
1019
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
1120
import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
1221
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
@@ -229,4 +238,65 @@ void testFunctionalSelectionExpressionWithNestedFieldWithAlias() {
229238
Params params = postgresQueryParser.getParamsBuilder().build();
230239
Assertions.assertEquals(0, params.getObjectParams().size());
231240
}
241+
242+
@Test
243+
void testAggregationExpression() {
244+
org.hypertrace.core.documentstore.query.Query query =
245+
org.hypertrace.core.documentstore.query.Query.builder()
246+
.setFilter(
247+
RelationalExpression.of(
248+
IdentifierExpression.of("price"), EQ, ConstantExpression.of(10)))
249+
.addSelection(IdentifierExpression.of("item"))
250+
.addSelection(
251+
AggregateExpression.of(AVG, IdentifierExpression.of("quantity")), "qty_avg")
252+
.addSelection(
253+
AggregateExpression.of(COUNT, IdentifierExpression.of("quantity")), "qty_count")
254+
.addSelection(
255+
AggregateExpression.of(DISTINCT_COUNT, IdentifierExpression.of("quantity")),
256+
"qty_distinct_count")
257+
.addSelection(
258+
AggregateExpression.of(SUM, IdentifierExpression.of("quantity")), "qty_sum")
259+
.addSelection(
260+
AggregateExpression.of(MIN, IdentifierExpression.of("quantity")), "qty_min")
261+
.addSelection(
262+
AggregateExpression.of(MAX, IdentifierExpression.of("quantity")), "qty_max")
263+
.addAggregation(IdentifierExpression.of("item"))
264+
.build();
265+
266+
PostgresQueryParser postgresQueryParser = new PostgresQueryParser(TEST_COLLECTION);
267+
String sql = postgresQueryParser.parse(query);
268+
Assertions.assertEquals(
269+
"SELECT document->'item' AS item, "
270+
+ "AVG( CAST (document->>'quantity' AS NUMERIC) ) AS qty_avg, "
271+
+ "COUNT( CAST (document->>'quantity' AS NUMERIC) ) AS qty_count, "
272+
+ "COUNT(DISTINCT CAST (document->>'quantity' AS NUMERIC) ) AS qty_distinct_count, "
273+
+ "SUM( CAST (document->>'quantity' AS NUMERIC) ) AS qty_sum, "
274+
+ "MIN( CAST (document->>'quantity' AS NUMERIC) ) AS qty_min, "
275+
+ "MAX( CAST (document->>'quantity' AS NUMERIC) ) AS qty_max "
276+
+ "FROM testCollection WHERE CAST (document->>'price' AS NUMERIC) = ? "
277+
+ "GROUP BY document->'item'",
278+
sql);
279+
280+
Params params = postgresQueryParser.getParamsBuilder().build();
281+
Assertions.assertEquals(1, params.getObjectParams().size());
282+
Assertions.assertEquals(10, params.getObjectParams().get(1));
283+
}
284+
285+
@Test
286+
void testAggregationExpressionDistinctCount() {
287+
org.hypertrace.core.documentstore.query.Query query =
288+
org.hypertrace.core.documentstore.query.Query.builder()
289+
.setFilter(
290+
RelationalExpression.of(
291+
IdentifierExpression.of("price"), EQ, ConstantExpression.of(10)))
292+
.addSelection(
293+
AggregateExpression.of(DISTINCT, IdentifierExpression.of("quantity")),
294+
"qty_distinct")
295+
.addAggregation(IdentifierExpression.of("item"))
296+
.build();
297+
298+
PostgresQueryParser postgresQueryParser = new PostgresQueryParser(TEST_COLLECTION);
299+
Assertions.assertThrows(
300+
UnsupportedOperationException.class, () -> postgresQueryParser.parse(query));
301+
}
232302
}

0 commit comments

Comments
 (0)