Skip to content

Commit 4daed08

Browse files
sarthak77Sarthak Singhal
sarthak77
and
Sarthak Singhal
authored
feat : Addition of API support for receiving attribute expressions (#126)
* added sup for attr expr in execution context * added sup for attr expr in projection transformation * added sup for attr expr in pinotbasedrequesthandler * refactored * added unit test for execution context * added unit test in pinot based rqst handler * changed function signature * resolved comments * nit * fixed unit tests * added template for new tests * added integ test with attr expr for existing queries * added functions for integ test * added new integ tests * nit * added contains key filter for selection Co-authored-by: Sarthak Singhal <sarthaksinghal@Sarthaks-MacBook-Pro.local>
1 parent 94d19d8 commit 4daed08

File tree

20 files changed

+777
-134
lines changed

20 files changed

+777
-134
lines changed

‎query-service-impl/src/main/java/org/hypertrace/core/query/service/ExecutionContext.java‎

Lines changed: 14 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
11
package org.hypertrace.core.query.service;
22

3+
import static org.hypertrace.core.query.service.QueryRequestUtil.getAlias;
4+
import static org.hypertrace.core.query.service.QueryRequestUtil.getLogicalColumnName;
5+
import static org.hypertrace.core.query.service.api.Expression.ValueCase.ATTRIBUTE_EXPRESSION;
6+
import static org.hypertrace.core.query.service.api.Expression.ValueCase.COLUMNIDENTIFIER;
7+
import static org.hypertrace.core.query.service.api.Expression.ValueCase.FUNCTION;
8+
39
import com.google.common.base.Supplier;
410
import com.google.common.base.Suppliers;
511
import java.time.Duration;
@@ -15,7 +21,6 @@
1521
import java.util.Optional;
1622
import java.util.Set;
1723
import java.util.concurrent.TimeUnit;
18-
import org.hypertrace.core.query.service.api.ColumnIdentifier;
1924
import org.hypertrace.core.query.service.api.ColumnMetadata;
2025
import org.hypertrace.core.query.service.api.Expression;
2126
import org.hypertrace.core.query.service.api.Expression.ValueCase;
@@ -137,26 +142,10 @@ private ColumnMetadata toColumnMetadata(Expression expression) {
137142
ValueCase valueCase = expression.getValueCase();
138143
switch (valueCase) {
139144
case COLUMNIDENTIFIER:
140-
ColumnIdentifier columnIdentifier = expression.getColumnIdentifier();
141-
String alias = columnIdentifier.getAlias();
142-
if (alias != null && alias.trim().length() > 0) {
143-
builder.setColumnName(alias);
144-
} else {
145-
builder.setColumnName(columnIdentifier.getColumnName());
146-
}
147-
builder.setValueType(ValueType.STRING);
148-
builder.setIsRepeated(false);
149-
break;
145+
case ATTRIBUTE_EXPRESSION:
150146
case FUNCTION:
151-
Function function = expression.getFunction();
152-
alias = function.getAlias();
153-
if (alias != null && alias.trim().length() > 0) {
154-
builder.setColumnName(alias);
155-
} else {
156-
// todo: handle recursive functions max(rollup(time,50)
157-
// workaround is to use alias for now
158-
builder.setColumnName(function.getFunctionName());
159-
}
147+
String alias = getAlias(expression).orElseThrow(IllegalArgumentException::new);
148+
builder.setColumnName(alias);
160149
builder.setValueType(ValueType.STRING);
161150
builder.setIsRepeated(false);
162151
break;
@@ -172,8 +161,10 @@ private void extractColumns(List<String> columns, Expression expression) {
172161
ValueCase valueCase = expression.getValueCase();
173162
switch (valueCase) {
174163
case COLUMNIDENTIFIER:
175-
ColumnIdentifier columnIdentifier = expression.getColumnIdentifier();
176-
columns.add(columnIdentifier.getColumnName());
164+
case ATTRIBUTE_EXPRESSION:
165+
String logicalColumnName =
166+
getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new);
167+
columns.add(logicalColumnName);
177168
break;
178169
case LITERAL:
179170
// no columns
@@ -233,7 +224,7 @@ private Optional<QueryTimeRange> buildQueryTimeRange(Filter filter, String timeF
233224
}
234225

235226
private boolean isMatchingFilter(Filter filter, String column, Collection<Operator> operators) {
236-
return column.equals(filter.getLhs().getColumnIdentifier().getColumnName())
227+
return getLogicalColumnName(filter.getLhs()).map(column::equals).orElse(false)
237228
&& (operators.stream()
238229
.anyMatch(operator -> Objects.equals(operator, filter.getOperator())));
239230
}

‎query-service-impl/src/main/java/org/hypertrace/core/query/service/QueryRequestUtil.java‎

Lines changed: 30 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22

33
import static org.hypertrace.core.query.service.api.Expression.ValueCase.ATTRIBUTE_EXPRESSION;
44
import static org.hypertrace.core.query.service.api.Expression.ValueCase.COLUMNIDENTIFIER;
5+
import static org.hypertrace.core.query.service.api.Expression.ValueCase.FUNCTION;
56

7+
import java.util.Optional;
68
import org.hypertrace.core.query.service.api.AttributeExpression;
79
import org.hypertrace.core.query.service.api.ColumnIdentifier;
810
import org.hypertrace.core.query.service.api.Expression;
@@ -124,19 +126,38 @@ public static boolean isDateTimeFunction(Expression expression) {
124126
&& expression.getFunction().getFunctionName().equals("dateTimeConvert");
125127
}
126128

127-
public static String getLogicalColumnName(Expression expression) {
129+
public static Optional<String> getLogicalColumnName(Expression expression) {
128130
switch (expression.getValueCase()) {
129131
case COLUMNIDENTIFIER:
130-
return expression.getColumnIdentifier().getColumnName();
132+
return Optional.of(expression.getColumnIdentifier().getColumnName());
131133
case ATTRIBUTE_EXPRESSION:
132-
return expression.getAttributeExpression().getAttributeId();
134+
return Optional.of(expression.getAttributeExpression().getAttributeId());
133135
default:
134-
throw new IllegalArgumentException(
135-
"Supports "
136-
+ ATTRIBUTE_EXPRESSION
137-
+ " and "
138-
+ COLUMNIDENTIFIER
139-
+ " expression type only");
136+
return Optional.empty();
137+
}
138+
}
139+
140+
public static Optional<String> getAlias(Expression expression) {
141+
switch (expression.getValueCase()) {
142+
case COLUMNIDENTIFIER:
143+
return Optional.of(
144+
expression.getColumnIdentifier().getAlias().isBlank()
145+
? getLogicalColumnName(expression).get()
146+
: expression.getColumnIdentifier().getAlias());
147+
case ATTRIBUTE_EXPRESSION:
148+
return Optional.of(
149+
expression.getAttributeExpression().getAlias().isBlank()
150+
? getLogicalColumnName(expression).get()
151+
: expression.getAttributeExpression().getAlias());
152+
case FUNCTION:
153+
// todo: handle recursive functions max(rollup(time,50)
154+
// workaround is to use alias for now
155+
return Optional.of(
156+
expression.getFunction().getAlias().isBlank()
157+
? expression.getFunction().getFunctionName()
158+
: expression.getFunction().getAlias());
159+
default:
160+
return Optional.empty();
140161
}
141162
}
142163
}

‎query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.hypertrace.core.query.service.pinot;
22

33
import static org.hypertrace.core.query.service.ConfigUtils.optionallyGet;
4+
import static org.hypertrace.core.query.service.QueryRequestUtil.getLogicalColumnName;
45

56
import com.google.common.base.Preconditions;
67
import com.google.common.base.Stopwatch;
@@ -225,10 +226,8 @@ && rhsHasLongValue(filter.getRhs())) {
225226
}
226227

227228
private boolean lhsIsStartTimeAttribute(Expression lhs) {
228-
return lhs.hasColumnIdentifier()
229-
&& startTimeAttributeName
230-
.map(attributeName -> attributeName.equals(lhs.getColumnIdentifier().getColumnName()))
231-
.orElse(false);
229+
return startTimeAttributeName.isPresent()
230+
&& startTimeAttributeName.equals(getLogicalColumnName(lhs));
232231
}
233232

234233
private boolean rhsHasLongValue(Expression rhs) {
@@ -245,7 +244,7 @@ private Set<String> getMatchingViewFilterColumns(
245244
// return it.
246245
if (filter.getChildFilterCount() == 0) {
247246
return doesSingleViewFilterMatchLeafQueryFilter(viewFilterMap, filter)
248-
? Set.of(filter.getLhs().getColumnIdentifier().getColumnName())
247+
? Set.of(getLogicalColumnName(filter.getLhs()).orElseThrow(IllegalArgumentException::new))
249248
: Set.of();
250249
} else {
251250
// 2. Internal filter node. Recursively get the matching nodes from children.
@@ -274,15 +273,13 @@ private Set<String> getMatchingViewFilterColumns(
274273
*/
275274
private boolean doesSingleViewFilterMatchLeafQueryFilter(
276275
Map<String, ViewColumnFilter> viewFilterMap, Filter queryFilter) {
277-
if (queryFilter.getLhs().getValueCase() != ValueCase.COLUMNIDENTIFIER) {
278-
return false;
279-
}
276+
280277
if (queryFilter.getOperator() != Operator.IN && queryFilter.getOperator() != Operator.EQ) {
281278
return false;
282279
}
283280

284-
String columnName = queryFilter.getLhs().getColumnIdentifier().getColumnName();
285-
ViewColumnFilter viewColumnFilter = viewFilterMap.get(columnName);
281+
ViewColumnFilter viewColumnFilter =
282+
viewFilterMap.get(getLogicalColumnName(queryFilter.getLhs()).orElse(null));
286283
if (viewColumnFilter == null) {
287284
return false;
288285
}
@@ -469,7 +466,8 @@ private Filter removeViewColumnFilter(
469466
private Filter rewriteLeafFilter(
470467
Filter queryFilter, Map<String, ViewColumnFilter> columnFilterMap) {
471468
ViewColumnFilter viewColumnFilter =
472-
columnFilterMap.get(queryFilter.getLhs().getColumnIdentifier().getColumnName());
469+
columnFilterMap.get(
470+
getLogicalColumnName(queryFilter.getLhs()).orElseThrow(IllegalArgumentException::new));
473471
// If the RHS of both the view filter and query filter match, return empty filter.
474472
if (viewColumnFilter != null && isEquals(viewColumnFilter.getValues(), queryFilter.getRhs())) {
475473
return Filter.getDefaultInstance();

‎query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/QueryRequestToPinotSQLConverter.java‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -211,7 +211,7 @@ private Expression handleValueConversionForLiteralExpression(Expression lhs, Exp
211211
return rhs;
212212
}
213213

214-
String lhsColumnName = getLogicalColumnName(lhs);
214+
String lhsColumnName = getLogicalColumnName(lhs).orElseThrow(IllegalArgumentException::new);
215215
try {
216216
Value value =
217217
DestinationColumnValueConverter.INSTANCE.convert(
@@ -272,7 +272,8 @@ private String convertExpressionToString(
272272
case COLUMNIDENTIFIER:
273273
// this takes care of the Map Type where it's split into 2 columns
274274
List<String> columnNames =
275-
viewDefinition.getPhysicalColumnNames(getLogicalColumnName(expression));
275+
viewDefinition.getPhysicalColumnNames(
276+
getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
276277
return joiner.join(columnNames);
277278
case ATTRIBUTE_EXPRESSION:
278279
if (isAttributeExpressionWithSubpath(expression)) {
@@ -292,7 +293,9 @@ private String convertExpressionToString(
292293
valCol);
293294
} else {
294295
// this takes care of the Map Type where it's split into 2 columns
295-
columnNames = viewDefinition.getPhysicalColumnNames(getLogicalColumnName(expression));
296+
columnNames =
297+
viewDefinition.getPhysicalColumnNames(
298+
getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
296299
return joiner.join(columnNames);
297300
}
298301
case LITERAL:
@@ -313,15 +316,19 @@ private String convertExpressionToString(
313316
}
314317

315318
private String convertExpressionToMapKeyColumn(Expression expression) {
316-
String col = viewDefinition.getKeyColumnNameForMap(getLogicalColumnName(expression));
319+
String col =
320+
viewDefinition.getKeyColumnNameForMap(
321+
getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
317322
if (col != null && col.length() > 0) {
318323
return col;
319324
}
320325
throw new IllegalArgumentException("operator supports multi value column only");
321326
}
322327

323328
private String convertExpressionToMapValueColumn(Expression expression) {
324-
String col = viewDefinition.getValueColumnNameForMap(getLogicalColumnName(expression));
329+
String col =
330+
viewDefinition.getValueColumnNameForMap(
331+
getLogicalColumnName(expression).orElseThrow(IllegalArgumentException::new));
325332
if (col != null && col.length() > 0) {
326333
return col;
327334
}

‎query-service-impl/src/main/java/org/hypertrace/core/query/service/projection/ProjectionTransformation.java‎

Lines changed: 60 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import java.util.List;
1818
import java.util.Optional;
1919
import java.util.stream.Collectors;
20+
import java.util.stream.Stream;
2021
import javax.inject.Inject;
2122
import org.hypertrace.core.attribute.service.cachingclient.CachingAttributeClient;
2223
import org.hypertrace.core.attribute.service.projection.AttributeProjection;
@@ -30,6 +31,7 @@
3031
import org.hypertrace.core.query.service.QueryFunctionConstants;
3132
import org.hypertrace.core.query.service.QueryRequestUtil;
3233
import org.hypertrace.core.query.service.QueryTransformation;
34+
import org.hypertrace.core.query.service.api.AttributeExpression;
3335
import org.hypertrace.core.query.service.api.ColumnIdentifier;
3436
import org.hypertrace.core.query.service.api.Expression;
3537
import org.hypertrace.core.query.service.api.Filter;
@@ -79,6 +81,8 @@ private Single<Expression> transformExpression(Expression expression) {
7981
switch (expression.getValueCase()) {
8082
case COLUMNIDENTIFIER:
8183
return this.transformColumnIdentifier(expression.getColumnIdentifier());
84+
case ATTRIBUTE_EXPRESSION:
85+
return this.transformAttributeExpression(expression.getAttributeExpression());
8286
case FUNCTION:
8387
return this.transformFunction(expression.getFunction())
8488
.map(expression.toBuilder()::setFunction)
@@ -96,10 +100,19 @@ private Single<Expression> transformExpression(Expression expression) {
96100

97101
private Single<Expression> transformColumnIdentifier(ColumnIdentifier columnIdentifier) {
98102
return this.projectAttributeIfPossible(columnIdentifier.getColumnName())
99-
.map(expression -> this.aliasToMatchOriginal(columnIdentifier, expression))
103+
.map(expression -> this.aliasToMatchOriginal(getOriginalKey(columnIdentifier), expression))
100104
.defaultIfEmpty(Expression.newBuilder().setColumnIdentifier(columnIdentifier).build());
101105
}
102106

107+
private Single<Expression> transformAttributeExpression(AttributeExpression attributeExpression) {
108+
return this.projectAttributeIfPossible(attributeExpression.getAttributeId())
109+
.map(
110+
expression ->
111+
this.aliasToMatchOriginal(getOriginalKey(attributeExpression), expression))
112+
.defaultIfEmpty(
113+
Expression.newBuilder().setAttributeExpression(attributeExpression).build());
114+
}
115+
103116
private Single<Function> transformFunction(Function function) {
104117
return this.transformExpressionList(function.getArgumentsList())
105118
.map(expressions -> function.toBuilder().clearArguments().addAllArguments(expressions))
@@ -267,15 +280,18 @@ private Single<String> convertOperator(ProjectionOperator operator) {
267280
}
268281
}
269282

270-
private Expression aliasToMatchOriginal(ColumnIdentifier original, Expression newExpression) {
271-
String originalKey =
272-
original.getAlias().isEmpty() ? original.getColumnName() : original.getAlias();
283+
private Expression aliasToMatchOriginal(String originalKey, Expression newExpression) {
273284
switch (newExpression.getValueCase()) {
274285
case COLUMNIDENTIFIER:
275286
return newExpression.toBuilder()
276287
.setColumnIdentifier(
277288
newExpression.getColumnIdentifier().toBuilder().setAlias(originalKey))
278289
.build();
290+
case ATTRIBUTE_EXPRESSION:
291+
return newExpression.toBuilder()
292+
.setAttributeExpression(
293+
newExpression.getAttributeExpression().toBuilder().setAlias(originalKey))
294+
.build();
279295
case FUNCTION:
280296
return newExpression.toBuilder()
281297
.setFunction(newExpression.getFunction().toBuilder().setAlias(originalKey))
@@ -329,7 +345,8 @@ private QueryRequest rebuildRequest(
329345
List<OrderByExpression> orderBys) {
330346

331347
QueryRequest.Builder builder = original.toBuilder();
332-
Filter updatedFilter = rebuildFilterForComplexAttributeExpression(originalFilter, orderBys);
348+
Filter updatedFilter =
349+
rebuildFilterForComplexAttributeExpression(originalFilter, orderBys, selections);
333350

334351
if (Filter.getDefaultInstance().equals(updatedFilter)) {
335352
builder.clearFilter();
@@ -350,16 +367,20 @@ private QueryRequest rebuildRequest(
350367
}
351368

352369
/*
353-
* We need the CONTAINS_KEY filter in all filters and order bys dealing with complex
370+
* We need the CONTAINS_KEY filter in all filters, selections and order bys dealing with complex
354371
* attribute expressions as Pinot gives error if particular key is absent. Rest all work fine.
355-
* To handle order bys, we add the corresponding filter at the top and 'AND' it with the main filter.
372+
* To handle order bys and selections, we add the corresponding filter at the top and 'AND' it with the main filter.
356373
* To handle filter, we modify each filter (say filter1) as : "CONTAINS_KEY AND filter1".
357374
*/
358375
private Filter rebuildFilterForComplexAttributeExpression(
359-
Filter originalFilter, List<OrderByExpression> orderBys) {
376+
Filter originalFilter, List<OrderByExpression> orderBys, List<Expression> selections) {
360377

361378
Filter updatedFilter = updateFilterForComplexAttributeExpressionFromFilter(originalFilter);
362-
List<Filter> filterList = createFilterForComplexAttributeExpressionFromOrderBy(orderBys);
379+
List<Filter> filterList =
380+
Stream.concat(
381+
createFilterForComplexAttributeExpressionFromOrderBy(orderBys),
382+
createFilterForComplexAttributeExpressionFromSelection(selections))
383+
.collect(Collectors.toList());
363384

364385
if (filterList.isEmpty()) {
365386
return updatedFilter;
@@ -408,13 +429,40 @@ private Filter updateFilterForComplexAttributeExpressionFromFilter(Filter origin
408429
}
409430
}
410431

411-
private List<Filter> createFilterForComplexAttributeExpressionFromOrderBy(
432+
private Stream<Filter> createFilterForComplexAttributeExpressionFromOrderBy(
412433
List<OrderByExpression> orderByExpressionList) {
413434
return orderByExpressionList.stream()
414435
.map(OrderByExpression::getExpression)
415436
.filter(QueryRequestUtil::isAttributeExpressionWithSubpath)
416437
.map(Expression::getAttributeExpression)
417-
.map(QueryRequestUtil::createContainsKeyFilter)
418-
.collect(Collectors.toList());
438+
.map(QueryRequestUtil::createContainsKeyFilter);
439+
}
440+
441+
private Stream<Filter> createFilterForComplexAttributeExpressionFromSelection(
442+
List<Expression> selections) {
443+
return selections.stream()
444+
.flatMap(this::getAnyAttributeExpression)
445+
.map(QueryRequestUtil::createContainsKeyFilter);
446+
}
447+
448+
private Stream<AttributeExpression> getAnyAttributeExpression(Expression selection) {
449+
if (selection.hasFunction()) {
450+
return selection.getFunction().getArgumentsList().stream()
451+
.flatMap(this::getAnyAttributeExpression);
452+
} else {
453+
return Stream.of(selection)
454+
.filter(QueryRequestUtil::isAttributeExpressionWithSubpath)
455+
.map(Expression::getAttributeExpression);
456+
}
457+
}
458+
459+
private String getOriginalKey(AttributeExpression attributeExpression) {
460+
String alias = attributeExpression.getAlias();
461+
return alias.isEmpty() ? attributeExpression.getAttributeId() : alias;
462+
}
463+
464+
private String getOriginalKey(ColumnIdentifier columnIdentifier) {
465+
String alias = columnIdentifier.getAlias();
466+
return alias.isEmpty() ? columnIdentifier.getColumnName() : alias;
419467
}
420468
}

‎query-service-impl/src/main/java/org/hypertrace/core/query/service/prometheus/FilterToPromqlConverter.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@ void convertFilterToString(
2828
}
2929
} else {
3030
if (QueryRequestUtil.isSimpleAttributeExpression(filter.getLhs())
31-
&& timeFilterColumn.equals(getLogicalColumnName(filter.getLhs()))) {
31+
&& timeFilterColumn.equals(
32+
getLogicalColumnName(filter.getLhs()).orElseThrow(IllegalArgumentException::new))) {
3233
return;
3334
}
3435
StringBuilder builder = new StringBuilder();

0 commit comments

Comments
 (0)