Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix : group by on functional expressions in mongo #200

Merged
merged 3 commits into from
Jul 2, 2024
Merged

fix : group by on functional expressions in mongo #200

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
93fdf17
fix : group by on functional expressions in mongo
RishabhB99 Jun 28, 2024
46737e1
addressed pr comments
RishabhB99 Jul 2, 2024
948d214
addressed pr comment
RishabhB99 Jul 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 50 additions & 0 deletions ...store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,11 @@
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.MAX;
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.MIN;
import static org.hypertrace.core.documentstore.expression.operators.AggregationOperator.SUM;
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.DIVIDE;
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.FLOOR;
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.LENGTH;
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.MULTIPLY;
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.SUBTRACT;
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.AND;
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.OR;
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.CONTAINS;
Expand Down Expand Up @@ -87,6 +90,7 @@
import org.hypertrace.core.documentstore.model.options.UpdateOptions;
import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue;
import org.hypertrace.core.documentstore.query.Aggregation;
import org.hypertrace.core.documentstore.query.Filter;
import org.hypertrace.core.documentstore.query.Pagination;
import org.hypertrace.core.documentstore.query.Query;
Expand Down Expand Up @@ -3312,6 +3316,52 @@ public void testNotExistsOperatorWithFindUsingBooleanRhs(String dataStoreName) t
testCountApi(dataStoreName, query, "query/not_exists_filter_response.json");
}

@ParameterizedTest
@ArgumentsSource(MongoProvider.class)
public void testMongoFunctionExpressionGroupBy(String dataStoreName) throws Exception {
Collection collection = getCollection(dataStoreName);

FunctionExpression functionExpression =
FunctionExpression.builder()
.operator(FLOOR)
.operand(
FunctionExpression.builder()
.operator(DIVIDE)
.operand(
FunctionExpression.builder()
.operator(SUBTRACT)
.operand(IdentifierExpression.of("price"))
.operand(ConstantExpression.of(5))
.build())
.operand(ConstantExpression.of(5))
.build())
.build();
List<SelectionSpec> selectionSpecs =
List.of(
SelectionSpec.of(functionExpression, "function"),
SelectionSpec.of(
AggregateExpression.of(COUNT, IdentifierExpression.of("function")),
"functionCount"));
Selection selection = Selection.builder().selectionSpecs(selectionSpecs).build();

Query query =
Query.builder()
.setSelection(selection)
.setAggregation(
Aggregation.builder().expression(IdentifierExpression.of("function")).build())
.setSort(
Sort.builder()
.sortingSpec(SortingSpec.of(IdentifierExpression.of("function"), ASC))
.build())
.build();

Iterator<Document> resultDocs = collection.aggregate(query);
assertDocsAndSizeEqualWithoutOrder(
dataStoreName, resultDocs, "query/function_expression_group_by_response.json", 3);

testCountApi(dataStoreName, query, "query/function_expression_group_by_response.json");
}

private static Collection getCollection(final String dataStoreName) {
return getCollection(dataStoreName, COLLECTION_NAME);
}
Expand Down
14 changes: 14 additions & 0 deletions ...ment-store/src/integrationTest/resources/query/function_expression_group_by_response.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
[
{
"function": 0.0,
"functionCount": 4
},
{
"function": 1.0,
"functionCount": 2
},
{
"function": 3.0,
"functionCount": 2
}
]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Empty line at the end.

19 changes: 19 additions & 0 deletions document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,16 @@
import java.util.Optional;
import java.util.Set;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;
import org.bson.json.JsonMode;
import org.bson.json.JsonWriterSettings;
import org.hypertrace.core.documentstore.Document;
import org.hypertrace.core.documentstore.JSONDocument;
import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
import org.hypertrace.core.documentstore.expression.type.GroupTypeExpression;
import org.hypertrace.core.documentstore.model.options.ReturnDocumentType;
import org.hypertrace.core.documentstore.query.SelectionSpec;

public final class MongoUtils {
public static final String FIELD_SEPARATOR = ".";
Expand Down Expand Up @@ -145,6 +150,20 @@ public static ReturnDocument getReturnDocument(final ReturnDocumentType returnDo
String.format("Unhandled return document type: %s", returnDocumentType)));
}

public static boolean isFunctionExpressionSelectionWithGroupBy(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Utils is only for very specific and common things. Can we move this method out of the utils?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is used at 2 places at MongoGroupTypeExpressionParser and MongoSelectTypeExpressionParser
so we'll need to define it at some common place, right?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It makes sense to put them in a common place, but doesn't feel like "Utils" is the right place. Either we could have a shared class (non-static) or have the definition in one place and access it in another.

final SelectionSpec selectionSpec, final List<String> groupByAliases) {
return selectionSpec.getAlias() != null
&& groupByAliases.contains(selectionSpec.getAlias())
&& selectionSpec.getExpression().getClass().equals(FunctionExpression.class);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please avoid checking class instances like this. Refer classes like IdentifierChecker, etc.

}

public static List<String> getGroupByAliases(final List<GroupTypeExpression> expressions) {
return expressions.stream()
.filter(expression -> expression.getClass().equals(IdentifierExpression.class))
.map(expression -> ((IdentifierExpression) expression).getName())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you avoid explicit instance of/class equality checks and casting by leveraging the visitor pattern?

.collect(Collectors.toUnmodifiableList());
}

private static ObjectNode wrapInLiteral(final ObjectNode objectNode) {
/* Wrapping the subDocument with $literal to be able to provide empty object "{}" as value
* Throws error otherwise if empty object is provided as value.
Expand Down
4 changes: 2 additions & 2 deletions ...store/src/main/java/org/hypertrace/core/documentstore/mongo/query/MongoQueryExecutor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
import static org.hypertrace.core.documentstore.mongo.query.MongoPaginationHelper.getSkipClause;
import static org.hypertrace.core.documentstore.mongo.query.parser.MongoFilterTypeExpressionParser.getFilter;
import static org.hypertrace.core.documentstore.mongo.query.parser.MongoFilterTypeExpressionParser.getFilterClause;
import static org.hypertrace.core.documentstore.mongo.query.parser.MongoGroupTypeExpressionParser.getGroupClause;
import static org.hypertrace.core.documentstore.mongo.query.parser.MongoNonProjectedSortTypeExpressionParser.getNonProjectedSortClause;
import static org.hypertrace.core.documentstore.mongo.query.parser.MongoSelectTypeExpressionParser.getProjectClause;
import static org.hypertrace.core.documentstore.mongo.query.parser.MongoSelectTypeExpressionParser.getSelections;
Expand Down Expand Up @@ -41,6 +40,7 @@
import org.hypertrace.core.documentstore.model.config.ConnectionConfig;
import org.hypertrace.core.documentstore.mongo.query.parser.AliasParser;
import org.hypertrace.core.documentstore.mongo.query.parser.MongoFromTypeExpressionParser;
import org.hypertrace.core.documentstore.mongo.query.parser.MongoGroupTypeExpressionParser;
import org.hypertrace.core.documentstore.mongo.query.transformer.MongoQueryTransformer;
import org.hypertrace.core.documentstore.parser.AggregateExpressionChecker;
import org.hypertrace.core.documentstore.parser.FunctionExpressionChecker;
Expand All @@ -58,7 +58,7 @@ public class MongoQueryExecutor {
List.of(
query -> singleton(getFilterClause(query, Query::getFilter)),
MongoFromTypeExpressionParser::getFromClauses,
query -> singleton(getGroupClause(query)),
MongoGroupTypeExpressionParser::getGroupClauses,
query -> singleton(getProjectClause(query)),
query -> singleton(getFilterClause(query, Query::getAggregationFilter)),
query -> singleton(getSortClause(query)),
Expand Down
3 changes: 2 additions & 1 deletion ...a/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,8 @@ Map<String, Object> parse(final FunctionExpression expression) {

SelectTypeExpressionVisitor parser =
new MongoIdentifierPrefixingParser(
new MongoIdentifierExpressionParser(new MongoConstantExpressionParser()));
new MongoIdentifierExpressionParser(
new MongoFunctionExpressionParser(new MongoConstantExpressionParser())));

if (numArgs == 1) {
Object value = expression.getOperands().get(0).accept(parser);
Expand Down
46 changes: 43 additions & 3 deletions .../org/hypertrace/core/documentstore/mongo/query/parser/MongoGroupTypeExpressionParser.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,16 @@
import static org.hypertrace.core.documentstore.mongo.MongoCollection.ID_KEY;
import static org.hypertrace.core.documentstore.mongo.MongoUtils.PREFIX;
import static org.hypertrace.core.documentstore.mongo.MongoUtils.encodeKey;
import static org.hypertrace.core.documentstore.mongo.MongoUtils.getGroupByAliases;
import static org.hypertrace.core.documentstore.mongo.MongoUtils.isFunctionExpressionSelectionWithGroupBy;

import com.mongodb.BasicDBObject;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
Expand All @@ -22,6 +26,7 @@
public final class MongoGroupTypeExpressionParser implements GroupTypeExpressionVisitor {

private static final String GROUP_CLAUSE = "$group";
private static final String ADD_FIELDS_CLAUSE = "$addFields";

@SuppressWarnings("unchecked")
@Override
Expand All @@ -41,10 +46,32 @@ public Map<String, Object> visit(final IdentifierExpression expression) {
return Map.of(key, PREFIX + identifier);
}

public static BasicDBObject getGroupClause(final Query query) {
public static List<BasicDBObject> getGroupClauses(final Query query) {
final List<SelectionSpec> selectionSpecs = query.getSelections();
final List<GroupTypeExpression> expressions = query.getAggregations();

final List<BasicDBObject> basicDBObjects = new ArrayList<>();

final List<SelectionSpec> functionExpressionSelectionWithGroupBys =
getFunctionExpressionSelectionWithGroupBys(selectionSpecs, expressions);

if (!functionExpressionSelectionWithGroupBys.isEmpty()) {
MongoSelectTypeExpressionParser parser =
new MongoIdentifierPrefixingParser(
new MongoIdentifierExpressionParser(new MongoFunctionExpressionParser()));
Map<String, Object> addFields =
functionExpressionSelectionWithGroupBys.stream()
.map(spec -> MongoGroupTypeExpressionParser.parse(parser, spec))
.reduce(
new LinkedHashMap<>(),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason for a linked hashmap? Is the order important here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no reason, i just copy pasted the reduce used below
don't think the order is important

(first, second) -> {
first.putAll(second);
return first;
});

basicDBObjects.add(new BasicDBObject(ADD_FIELDS_CLAUSE, addFields));
}

MongoGroupTypeExpressionParser parser = new MongoGroupTypeExpressionParser();
Map<String, Object> groupExp;

Expand Down Expand Up @@ -82,11 +109,13 @@ public static BasicDBObject getGroupClause(final Query query) {
});

if (MapUtils.isEmpty(definition) && CollectionUtils.isEmpty(expressions)) {
return new BasicDBObject();
return basicDBObjects;
}

definition.putAll(groupExp);
return new BasicDBObject(GROUP_CLAUSE, definition);

basicDBObjects.add(new BasicDBObject(GROUP_CLAUSE, definition));
return basicDBObjects;
}

private static Map<String, Object> parse(
Expand All @@ -99,4 +128,15 @@ private Map<String, Object> parse(final GroupTypeExpression expression) {
MongoGroupTypeExpressionParser parser = new MongoGroupTypeExpressionParser();
return expression.accept(parser);
}

private static List<SelectionSpec> getFunctionExpressionSelectionWithGroupBys(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is this getting used?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

document-store/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoGroupTypeExpressionParser.java

Line 56 in 93fdf17

getFunctionExpressionSelectionWithGroupBys(selectionSpecs, expressions);

final List<SelectionSpec> selectionSpecs, final List<GroupTypeExpression> expressions) {
List<String> groupByAliases = getGroupByAliases(expressions);

return selectionSpecs.stream()
.filter(
selectionSpec ->
isFunctionExpressionSelectionWithGroupBy(selectionSpec, groupByAliases))
.collect(Collectors.toUnmodifiableList());
}
}
15 changes: 15 additions & 0 deletions ...org/hypertrace/core/documentstore/mongo/query/parser/MongoSelectTypeExpressionParser.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
package org.hypertrace.core.documentstore.mongo.query.parser;

import static java.util.stream.Collectors.toMap;
import static org.hypertrace.core.documentstore.mongo.MongoCollection.ID_KEY;
import static org.hypertrace.core.documentstore.mongo.MongoUtils.getGroupByAliases;
import static org.hypertrace.core.documentstore.mongo.MongoUtils.isFunctionExpressionSelectionWithGroupBy;

import com.google.common.base.Joiner;
import com.mongodb.BasicDBObject;
import java.util.List;
import java.util.Map;
Expand All @@ -20,6 +24,8 @@ public abstract class MongoSelectTypeExpressionParser implements SelectTypeExpre

protected final MongoSelectTypeExpressionParser baseParser;

private static final Joiner DOT_JOINER = Joiner.on(".");

protected MongoSelectTypeExpressionParser() {
this(MongoUnsupportedSelectTypeExpressionParser.INSTANCE);
}
Expand Down Expand Up @@ -59,8 +65,17 @@ public static BasicDBObject getSelections(final Query query) {
new MongoIdentifierPrefixingParser(
new MongoIdentifierExpressionParser(new MongoFunctionExpressionParser()));

List<String> groupByAliases = getGroupByAliases(query.getAggregations());

Map<String, Object> projectionMap =
selectionSpecs.stream()
.map(
spec ->
isFunctionExpressionSelectionWithGroupBy(spec, groupByAliases)
? SelectionSpec.of(
IdentifierExpression.of(DOT_JOINER.join(ID_KEY, spec.getAlias())),
spec.getAlias())
: spec)
.map(spec -> MongoSelectTypeExpressionParser.parse(parser, spec))
.flatMap(map -> map.entrySet().stream())
.collect(
Expand Down
Loading