Skip to content

Commit 01ad553

Browse files
anshul98ks123Anshul Singh
anshul98ks123
and
Anshul Singh
authored
[persistence] Optimize anomalies list API with combined query (#1800)
* [persistence] Optimize anomalies list API with combined query * handle datetime in where clause and use JOIN instead of IN * Add daoFilter validation, tests for getV2 and boolean handling in where predicate * cleanup * Add missing final * address feedback * cleanup * Update method description of getV2 --------- Co-authored-by: Anshul Singh <anshul.singh@anshuls-macbook-pro-1.wyvern-sun.ts.net>
1 parent c346f16 commit 01ad553

File tree

6 files changed

+270
-2
lines changed

6 files changed

+270
-2
lines changed

‎thirdeye-persistence/src/main/java/ai/startree/thirdeye/datalayer/DatabaseOrm.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,11 @@ public <E extends AbstractEntity> String getIdColumnName(final Class<E> clazz) {
148148
return AbstractIndexEntity.class.isAssignableFrom(clazz) ? "baseId" : "id";
149149
}
150150

151+
public <E extends AbstractEntity> String getIdColumnSQLName(final Class<E> clazz) {
152+
final String idColName = getIdColumnName(clazz);
153+
return sqlQueryBuilder.getColumnSQLName(clazz, idColName);
154+
}
155+
151156
public Integer delete(final Predicate predicate,
152157
final Class<? extends AbstractEntity> entityClass, final Connection connection)
153158
throws Exception {
@@ -203,4 +208,10 @@ public <E extends AbstractEntity> List<E> runSQL(
203208
sample.stop(dbCrudTimerOfRead);
204209
}
205210
}
211+
212+
public <E extends AbstractEntity> String generateMatchingIdsQuery(
213+
final Predicate predicate, final Long limit, final Long offset, final Class<E> clazz) {
214+
return sqlQueryBuilder.createFindColumnByParamsStatementWithLimitQuery(
215+
clazz, getIdColumnName(clazz), predicate, limit, offset);
216+
}
206217
}

‎thirdeye-persistence/src/main/java/ai/startree/thirdeye/datalayer/bao/AbstractManagerImpl.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,10 @@ public List<E> filter(final DaoFilter daoFilter) {
150150
return genericPojoDao.get(daoFilter.setBeanClass(dtoClass));
151151
}
152152

153+
public List<E> filterV2(final DaoFilter daoFilter) {
154+
return genericPojoDao.getV2(daoFilter.setBeanClass(dtoClass));
155+
}
156+
153157
@Override
154158
public long count() {
155159
return genericPojoDao.count(dtoClass);

‎thirdeye-persistence/src/main/java/ai/startree/thirdeye/datalayer/bao/AnomalyManagerImpl.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ public List<AnomalyDTO> findAll() {
195195

196196
@Override
197197
public List<AnomalyDTO> filter(final DaoFilter daoFilter) {
198-
final List<AnomalyDTO> anomalies = super.filter(daoFilter);
198+
final List<AnomalyDTO> anomalies = super.filterV2(daoFilter);
199199
// FIXME CYRIL this filter is only decorating with feedback - while some others decorate with feedback and children
200200
return decorateWithFeedback(anomalies);
201201
}

‎thirdeye-persistence/src/main/java/ai/startree/thirdeye/datalayer/dao/GenericPojoDao.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
import java.lang.reflect.Field;
4444
import java.sql.Timestamp;
4545
import java.util.ArrayList;
46+
import java.util.Collections;
4647
import java.util.List;
4748
import java.util.Set;
4849
import org.apache.commons.collections4.CollectionUtils;
@@ -399,4 +400,58 @@ public <E extends AbstractDTO> int deleteByPredicate(final Predicate predicate,
399400
new DaoFilter().setPredicate(predicate).setBeanClass(pojoClass));
400401
return delete(idsToDelete, pojoClass);
401402
}
403+
404+
/**
405+
* Use this method when you want to fetch entities out a subset of the entities based on predicates,
406+
* limits, offsets, etc.
407+
*
408+
* This method is an optimized version of get(final DaoFilter daoFilter)
409+
* where we make a single combined query of the form
410+
*
411+
* select generic_json_entity.* from generic_json_entity
412+
* JOIN (select base_id from index_class where CONDITIONS)
413+
* subquery ON generic_json_entity.id = subquery.base_id
414+
*
415+
* @param daoFilter required filters to fetch the result.
416+
*/
417+
public <E extends AbstractDTO> List<E> getV2(final DaoFilter daoFilter) {
418+
final Class<? extends AbstractIndexEntity> indexClass = BEAN_INDEX_MAP.get(
419+
daoFilter.getBeanClass());
420+
validate(daoFilter);
421+
422+
final String matchingIdsQuery = databaseOrm.generateMatchingIdsQuery(
423+
daoFilter.getPredicate(),
424+
daoFilter.getLimit(),
425+
daoFilter.getOffset(),
426+
indexClass);
427+
final String parameterizedSQL = String.format("""
428+
JOIN (
429+
%s
430+
) subquery ON generic_json_entity.%s = subquery.%s
431+
""",
432+
matchingIdsQuery,
433+
databaseOrm.getIdColumnSQLName(GenericJsonEntity.class),
434+
databaseOrm.getIdColumnSQLName(indexClass));
435+
436+
try {
437+
final List<GenericJsonEntity> entities = databaseClient.executeTransaction(
438+
(connection) -> databaseOrm.runSQL(
439+
parameterizedSQL,
440+
Collections.emptyMap(),
441+
GenericJsonEntity.class,
442+
connection));
443+
final List<E> results = new ArrayList<>();
444+
if (CollectionUtils.isNotEmpty(entities)) {
445+
for (final GenericJsonEntity entity : entities) {
446+
final E e = toDto(entity, (Class<E>) daoFilter.getBeanClass());
447+
results.add(e);
448+
}
449+
}
450+
return results;
451+
} catch (Exception e) {
452+
LOG.error(e.getMessage(), e);
453+
// TODO ANSHUL design - surface exception ?
454+
return Collections.emptyList();
455+
}
456+
}
402457
}

‎thirdeye-persistence/src/main/java/ai/startree/thirdeye/datalayer/util/SqlQueryBuilder.java‎

Lines changed: 175 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@
2828
import java.sql.PreparedStatement;
2929
import java.sql.Statement;
3030
import java.sql.Types;
31+
import java.time.LocalDateTime;
32+
import java.time.format.DateTimeFormatter;
33+
import java.time.format.DateTimeParseException;
3134
import java.util.ArrayList;
3235
import java.util.Date;
3336
import java.util.HashMap;
@@ -243,6 +246,52 @@ public PreparedStatement createFindByParamsStatementWithLimit(final Connection c
243246
return prepareStatement;
244247
}
245248

249+
public String getColumnSQLName(final Class<? extends AbstractEntity> entityClass,
250+
final String column) {
251+
final String tableName = entityMappingHolder.tableToEntityNameMap.inverse()
252+
.get(entityClass.getSimpleName());
253+
final BiMap<String, String> entityNameToDBNameMapping =
254+
entityMappingHolder.columnMappingPerTable.get(tableName).inverse();
255+
256+
final String columnName = entityNameToDBNameMapping.get(column);
257+
checkNotNull(columnName, String
258+
.format("Found field '%s' but expected %s", column,
259+
entityNameToDBNameMapping.keySet()));
260+
261+
return columnName;
262+
}
263+
264+
public String createFindColumnByParamsStatementWithLimitQuery(
265+
final Class<? extends AbstractEntity> entityClass, final String column,
266+
final Predicate predicate, final Long limit, final Long offset) {
267+
final String tableName = entityMappingHolder.tableToEntityNameMap.inverse()
268+
.get(entityClass.getSimpleName());
269+
final BiMap<String, String> entityNameToDBNameMapping =
270+
entityMappingHolder.columnMappingPerTable.get(tableName).inverse();
271+
272+
final String columnName = entityNameToDBNameMapping.get(column);
273+
checkNotNull(columnName, String
274+
.format("Found field '%s' but expected %s", column,
275+
entityNameToDBNameMapping.keySet()));
276+
277+
final StringBuilder sqlBuilder = new StringBuilder(String.format(
278+
"SELECT %s FROM %s", columnName, tableName));
279+
280+
if(predicate != null) {
281+
final StringBuilder whereClause = new StringBuilder(" WHERE ");
282+
createWhereClause(entityNameToDBNameMapping, predicate, whereClause);
283+
sqlBuilder.append(whereClause);
284+
}
285+
if (limit != null) {
286+
sqlBuilder.append(" LIMIT ").append(limit);
287+
}
288+
if (offset != null) {
289+
sqlBuilder.append(" OFFSET ").append(offset);
290+
}
291+
return sqlBuilder.toString();
292+
}
293+
294+
246295
public PreparedStatement createCountStatement(final Connection connection, final @Nullable Predicate predicate,
247296
final Class<? extends AbstractEntity> entityClass) throws Exception {
248297
final String tableName =
@@ -363,12 +412,137 @@ private void generateWhereClause(final BiMap<String, String> entityNameToDBNameM
363412
}
364413
}
365414

415+
private void createWhereClause(final BiMap<String, String> entityNameToDBNameMapping,
416+
final Predicate predicate, final StringBuilder whereClause) {
417+
String columnName = null;
418+
419+
if (predicate.getLhs() != null) {
420+
columnName = entityNameToDBNameMapping.get(predicate.getLhs());
421+
checkNotNull(columnName, String
422+
.format("Found field '%s' but expected %s", predicate.getLhs(),
423+
entityNameToDBNameMapping.keySet()));
424+
}
425+
426+
switch (predicate.getOper()) {
427+
case AND:
428+
case OR:
429+
whereClause.append("(");
430+
String delim = "";
431+
for (final Predicate childPredicate : predicate.getChildPredicates()) {
432+
whereClause.append(delim);
433+
createWhereClause(entityNameToDBNameMapping, childPredicate, whereClause);
434+
delim = " " + predicate.getOper().toString() + " ";
435+
}
436+
whereClause.append(")");
437+
break;
438+
case EQ:
439+
if (predicate.getRhs() == null) {
440+
whereClause.append(columnName).append(" IS NULL ");
441+
} else {
442+
// duplicated code with NEQ and LIKE/GT/GE/... - ok for the moment, this needs to be migrated to JOOQ anyway
443+
whereClause.append(columnName).append(" ").append(predicate.getOper().toString())
444+
.append(" ").append(getPredicateValStr(predicate.getRhs()));
445+
}
446+
break;
447+
case NEQ:
448+
if (predicate.getRhs() == null) {
449+
whereClause.append(columnName).append(" IS NOT NULL ");
450+
} else {
451+
whereClause.append(columnName).append(" ").append(predicate.getOper().toString())
452+
.append(" ").append(getPredicateValStr(predicate.getRhs()));
453+
}
454+
break;
455+
case LIKE:
456+
case GT:
457+
case LT:
458+
case LE:
459+
case GE:
460+
whereClause.append(columnName).append(" ").append(predicate.getOper().toString())
461+
.append(" ").append(getPredicateValStr(predicate.getRhs()));
462+
break;
463+
case IN:
464+
Object rhs = predicate.getRhs();
465+
if (rhs != null) {
466+
if (!rhs.getClass().isArray()) {
467+
rhs = rhs.toString().split(",");
468+
}
469+
whereClause.append(columnName).append(" ").append(Predicate.OPER.IN)
470+
.append("(");
471+
delim = "";
472+
final int length = Array.getLength(rhs);
473+
if (length > 0) {
474+
for (int i = 0; i < length; i++) {
475+
whereClause.append(delim).append(getPredicateValStr(Array.get(rhs, i)));
476+
delim = ",";
477+
}
478+
} else {
479+
whereClause.append("null");
480+
}
481+
whereClause.append(")");
482+
}
483+
break;
484+
case BETWEEN:
485+
final ImmutablePair<Object, Object> pair = (ImmutablePair<Object, Object>) predicate.getRhs();
486+
whereClause.append(columnName).append(predicate.getOper().toString())
487+
.append(getPredicateValStr(pair.getLeft()))
488+
.append(" AND ")
489+
.append(getPredicateValStr(pair.getRight()));
490+
break;
491+
default:
492+
throw new RuntimeException("Unsupported predicate type:" + predicate.getOper());
493+
}
494+
}
495+
496+
private String getPredicateValStr(Object val) {
497+
if (checkIfValidBoolean(val)) {
498+
return val.toString();
499+
}
500+
else if (val instanceof String || checkIfValidDateTime(val)) {
501+
return "'" + val + "'";
502+
} else {
503+
return val.toString();
504+
}
505+
}
506+
507+
private boolean checkIfValidDateTime(Object val) {
508+
// List of possible patterns
509+
final String[] patterns = {
510+
"yyyy-MM-dd HH:mm:ss.SSS", // e.g. 2025-02-17 17:45:06.493
511+
"yyyy-MM-dd HH:mm:ss.SS", // e.g. 2020-02-17 23:30:00.28
512+
"yyyy-MM-dd HH:mm:ss.S", // e.g. 2020-02-17 23:30:00.0
513+
"yyyy-MM-dd HH:mm:ss", // e.g. 2025-02-17 17:45:06
514+
"yyyy-MM-dd'T'HH:mm:ss.SSS", // e.g. 2025-02-17T17:45:06.493
515+
"yyyy-MM-dd'T'HH:mm:ss.SS", // e.g. 2025-02-17T17:45:06.49
516+
"yyyy-MM-dd'T'HH:mm:ss.S", // e.g. 2020-02-17T23:30:00.0
517+
"yyyy-MM-dd'T'HH:mm:ss", // e.g. 2025-02-17T17:45:06
518+
"yyyy/MM/dd HH:mm:ss", // e.g. 2025/02/17 17:45:06
519+
"yyyy/MM/dd", // e.g. 2025/02/17
520+
"MM/dd/yyyy HH:mm:ss" // e.g. 02/17/2025 17:45:06
521+
};
522+
523+
for (String pattern : patterns) {
524+
try {
525+
final DateTimeFormatter formatter = DateTimeFormatter.ofPattern(pattern);
526+
LocalDateTime.parse(val.toString(), formatter);
527+
return true; // If parsing is successful with any pattern
528+
} catch (DateTimeParseException e) {
529+
// If parsing fails, continue with next pattern
530+
}
531+
}
532+
533+
return false;
534+
}
535+
536+
private boolean checkIfValidBoolean(Object val) {
537+
return "true".equalsIgnoreCase(val.toString()) || "false".equalsIgnoreCase(val.toString());
538+
}
539+
366540
public PreparedStatement createStatementFromSQL(final Connection connection, String parameterizedSQL,
367541
final Map<String, Object> parameterMap, final Class<? extends AbstractEntity> entityClass)
368542
throws Exception {
369543
final String tableName =
370544
entityMappingHolder.tableToEntityNameMap.inverse().get(entityClass.getSimpleName());
371-
parameterizedSQL = "select * from " + tableName + " " + parameterizedSQL;
545+
parameterizedSQL = "select " + tableName + ".* from " + tableName + " " + parameterizedSQL;
372546
parameterizedSQL = parameterizedSQL.replace(entityClass.getSimpleName(), tableName);
373547
final StringBuilder psSql = new StringBuilder();
374548
final List<String> paramNames = new ArrayList<>();

‎thirdeye-persistence/src/test/java/ai/startree/thirdeye/datalayer/dao/TestGenericPojoDao.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,10 @@ public void updateEntityTest() {
100100
new DaoFilter().setPredicate(Predicate.EQ(TYPE, TEST_TYPES.get(0)))
101101
.setBeanClass(DataSourceDTO.class));
102102
assertThat(dtos.size()).isEqualTo(2);
103+
final List<DataSourceDTO> dtosV2 = dao.getV2(
104+
new DaoFilter().setPredicate(Predicate.EQ(TYPE, TEST_TYPES.get(0)))
105+
.setBeanClass(DataSourceDTO.class));
106+
assertThat(dtosV2.size()).isEqualTo(2);
103107
final DataSourceDTO dto = dtos.get(0);
104108
dto.setType(TEST_TYPES.get(1));
105109
final long idBeforeUpdate = dto.getId();
@@ -109,6 +113,10 @@ public void updateEntityTest() {
109113
final List<DataSourceDTO> dtoAfterUpdate = dao.get(
110114
new DaoFilter().setPredicate(Predicate.EQ(TYPE, TEST_TYPES.get(1)))
111115
.setBeanClass(DataSourceDTO.class));
116+
final List<DataSourceDTO> dtoV2AfterUpdate = dao.getV2(
117+
new DaoFilter().setPredicate(Predicate.EQ(TYPE, TEST_TYPES.get(1)))
118+
.setBeanClass(DataSourceDTO.class));
119+
assertThat(dtoAfterUpdate).isEqualTo(dtoV2AfterUpdate);
112120
assertThat(dtoAfterUpdate.size()).isEqualTo(1);
113121
assertThat(dtoAfterUpdate.get(0).getId()).isEqualTo(idBeforeUpdate);
114122
}
@@ -129,6 +137,8 @@ public void filterWithLimitTest() {
129137
.setLimit(limit)
130138
.setBeanClass(AnomalyDTO.class);
131139
final List<AnomalyDTO> anomalies = dao.get(filter);
140+
final List<AnomalyDTO> anomaliesV2 = dao.get(filter);
141+
assertThat(anomalies).isEqualTo(anomaliesV2);
132142
assertThat(anomalies).isNotNull();
133143
assertThat(anomalies.size()).isEqualTo(limit);
134144
}
@@ -143,6 +153,8 @@ public void filterWithLimitAndOffsetTest() {
143153
.setOffset(offset)
144154
.setBeanClass(AnomalyDTO.class);
145155
final List<AnomalyDTO> anomalies = dao.get(filter);
156+
final List<AnomalyDTO> anomaliesV2 = dao.getV2(filter);
157+
assertThat(anomalies).isEqualTo(anomaliesV2);
146158
assertThat(anomalies).isNotNull();
147159
assertThat(anomalies.size()).isEqualTo(TOTAL_ANOMALIES - offset);
148160
}
@@ -156,6 +168,9 @@ public void filterWithOffsetWithoutLimitTest() {
156168
assertThatThrownBy(() -> dao.get(filter))
157169
.isInstanceOf(IllegalArgumentException.class)
158170
.hasMessage(ERR_OFFSET_WITHOUT_LIMIT.getMessage());
171+
assertThatThrownBy(() -> dao.getV2(filter))
172+
.isInstanceOf(IllegalArgumentException.class)
173+
.hasMessage(ERR_OFFSET_WITHOUT_LIMIT.getMessage());
159174
}
160175

161176
@Test
@@ -166,6 +181,9 @@ public void testNegativeLimitValue() {
166181
assertThatThrownBy(() -> dao.get(filter))
167182
.isInstanceOf(IllegalArgumentException.class)
168183
.hasMessage(ERR_NEGATIVE_LIMIT_VALUE.getMessage());
184+
assertThatThrownBy(() -> dao.getV2(filter))
185+
.isInstanceOf(IllegalArgumentException.class)
186+
.hasMessage(ERR_NEGATIVE_LIMIT_VALUE.getMessage());
169187
}
170188

171189
@Test
@@ -177,6 +195,9 @@ public void testNegativeOffsetValue() {
177195
assertThatThrownBy(() -> dao.get(filter))
178196
.isInstanceOf(IllegalArgumentException.class)
179197
.hasMessage(ERR_NEGATIVE_OFFSET_VALUE.getMessage());
198+
assertThatThrownBy(() -> dao.getV2(filter))
199+
.isInstanceOf(IllegalArgumentException.class)
200+
.hasMessage(ERR_NEGATIVE_OFFSET_VALUE.getMessage());
180201
}
181202

182203
@Test
@@ -192,7 +213,10 @@ public void pageBoundariesTest() {
192213
while (offset < TOTAL_ANOMALIES) {
193214
filter.setOffset(offset);
194215
final List<AnomalyDTO> anomalies = dao.get(filter);
216+
final List<AnomalyDTO> anomaliesV2 = dao.getV2(filter);
195217
assertThat(anomalies).isNotNull();
218+
assertThat(anomaliesV2).isNotNull();
219+
assertThat(anomalies).isEqualTo(anomaliesV2);
196220
entryCount += anomalies.size();
197221
// limit -> for all but the last page
198222
// TOTAL_ANOMALIES - offset -> for the last page

0 commit comments

Comments
 (0)