Skip to content

Commit 67ee096

Browse files
malwaregarrymalwaregarry
committed
[Databend] Implement NoREC expression generator interface
1 parent ef59534 commit 67ee096

6 files changed

Lines changed: 124 additions & 13 deletions

File tree

‎src/sqlancer/databend/ast/DatabendExpression.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package sqlancer.databend.ast;
22

3+
import sqlancer.common.ast.newast.Expression;
4+
import sqlancer.databend.DatabendSchema.DatabendColumn;
35
import sqlancer.databend.DatabendSchema.DatabendDataType;
46

5-
public interface DatabendExpression {
7+
public interface DatabendExpression extends Expression<DatabendColumn> {
68

79
default DatabendDataType getExpectedType() {
810
return null;

‎src/sqlancer/databend/ast/DatabendJoin.java‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,20 @@
44
import java.util.List;
55

66
import sqlancer.Randomly;
7+
import sqlancer.common.ast.newast.Join;
78
import sqlancer.common.ast.newast.TableReferenceNode;
89
import sqlancer.databend.DatabendProvider.DatabendGlobalState;
910
import sqlancer.databend.DatabendSchema;
1011
import sqlancer.databend.DatabendSchema.DatabendColumn;
1112
import sqlancer.databend.DatabendSchema.DatabendTable;
1213
import sqlancer.databend.gen.DatabendNewExpressionGenerator;
1314

14-
public class DatabendJoin implements DatabendExpression {
15+
public class DatabendJoin implements DatabendExpression, Join<DatabendExpression, DatabendTable, DatabendColumn> {
1516

1617
private final DatabendTableReference leftTable;
1718
private final DatabendTableReference rightTable;
1819
private final JoinType joinType;
19-
private final DatabendExpression onCondition;
20+
private DatabendExpression onCondition;
2021
private OuterType outerType;
2122

2223
public enum JoinType {
@@ -67,9 +68,8 @@ public OuterType getOuterType() {
6768
return outerType;
6869
}
6970

70-
public static List<DatabendExpression> getJoins(List<DatabendTableReference> tableList,
71-
DatabendGlobalState globalState) {
72-
List<DatabendExpression> joinExpressions = new ArrayList<>();
71+
public static List<DatabendJoin> getJoins(List<DatabendTableReference> tableList, DatabendGlobalState globalState) {
72+
List<DatabendJoin> joinExpressions = new ArrayList<>();
7373
while (tableList.size() >= 2 && Randomly.getBooleanWithRatherLowProbability()) {
7474
DatabendTableReference leftTable = tableList.remove(0);
7575
DatabendTableReference rightTable = tableList.remove(0);
@@ -116,11 +116,15 @@ public static DatabendJoin createInnerJoin(DatabendTableReference left, Databend
116116
return new DatabendJoin(left, right, JoinType.INNER, predicate);
117117
}
118118

119-
public static DatabendExpression createNaturalJoin(DatabendTableReference left, DatabendTableReference right,
119+
public static DatabendJoin createNaturalJoin(DatabendTableReference left, DatabendTableReference right,
120120
OuterType naturalJoinType) {
121121
DatabendJoin join = new DatabendJoin(left, right, JoinType.NATURAL, null);
122122
join.setOuterType(naturalJoinType);
123123
return join;
124124
}
125125

126+
@Override
127+
public void setOnClause(DatabendExpression onClause) {
128+
onCondition = onClause;
129+
}
126130
}

‎src/sqlancer/databend/ast/DatabendSelect.java‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,16 @@
11
package sqlancer.databend.ast;
22

3+
import java.util.List;
4+
import java.util.stream.Collectors;
5+
36
import sqlancer.common.ast.SelectBase;
7+
import sqlancer.common.ast.newast.Select;
8+
import sqlancer.databend.DatabendSchema.DatabendColumn;
9+
import sqlancer.databend.DatabendSchema.DatabendTable;
10+
import sqlancer.databend.DatabendToStringVisitor;
411

5-
public class DatabendSelect extends SelectBase<DatabendExpression> implements DatabendExpression {
12+
public class DatabendSelect extends SelectBase<DatabendExpression>
13+
implements DatabendExpression, Select<DatabendJoin, DatabendExpression, DatabendTable, DatabendColumn> {
614

715
private boolean isDistinct;
816

@@ -14,4 +22,20 @@ public boolean isDistinct() {
1422
return isDistinct;
1523
}
1624

25+
@Override
26+
public void setJoinClauses(List<DatabendJoin> joinStatements) {
27+
List<DatabendExpression> expressions = joinStatements.stream().map(e -> (DatabendExpression) e)
28+
.collect(Collectors.toList());
29+
setJoinList(expressions);
30+
}
31+
32+
@Override
33+
public List<DatabendJoin> getJoinClauses() {
34+
return getJoinList().stream().map(e -> (DatabendJoin) e).collect(Collectors.toList());
35+
}
36+
37+
@Override
38+
public String asString() {
39+
return DatabendToStringVisitor.asString(this);
40+
}
1741
}

‎src/sqlancer/databend/gen/DatabendNewExpressionGenerator.java‎

Lines changed: 83 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,17 @@
99

1010
import sqlancer.Randomly;
1111
import sqlancer.common.ast.newast.NewOrderingTerm;
12+
import sqlancer.common.gen.NoRECGenerator;
1213
import sqlancer.common.gen.TypedExpressionGenerator;
14+
import sqlancer.common.schema.AbstractTables;
1315
import sqlancer.databend.DatabendBugs;
1416
import sqlancer.databend.DatabendProvider.DatabendGlobalState;
1517
import sqlancer.databend.DatabendSchema.DatabendColumn;
18+
import sqlancer.databend.DatabendSchema.DatabendCompositeDataType;
1619
import sqlancer.databend.DatabendSchema.DatabendDataType;
1720
import sqlancer.databend.DatabendSchema.DatabendRowValue;
21+
import sqlancer.databend.DatabendSchema.DatabendTable;
22+
import sqlancer.databend.DatabendToStringVisitor;
1823
import sqlancer.databend.ast.DatabendAggregateOperation;
1924
import sqlancer.databend.ast.DatabendAggregateOperation.DatabendAggregateFunction;
2025
import sqlancer.databend.ast.DatabendBetweenOperation;
@@ -24,21 +29,29 @@
2429
import sqlancer.databend.ast.DatabendBinaryComparisonOperation.DatabendBinaryComparisonOperator;
2530
import sqlancer.databend.ast.DatabendBinaryLogicalOperation;
2631
import sqlancer.databend.ast.DatabendBinaryLogicalOperation.DatabendBinaryLogicalOperator;
32+
import sqlancer.databend.ast.DatabendCastOperation;
33+
import sqlancer.databend.ast.DatabendColumnReference;
2734
import sqlancer.databend.ast.DatabendColumnValue;
2835
import sqlancer.databend.ast.DatabendConstant;
2936
import sqlancer.databend.ast.DatabendExpression;
3037
import sqlancer.databend.ast.DatabendInOperation;
38+
import sqlancer.databend.ast.DatabendJoin;
3139
import sqlancer.databend.ast.DatabendLikeOperation;
3240
import sqlancer.databend.ast.DatabendOrderByTerm;
41+
import sqlancer.databend.ast.DatabendPostFixText;
42+
import sqlancer.databend.ast.DatabendSelect;
43+
import sqlancer.databend.ast.DatabendTableReference;
3344
import sqlancer.databend.ast.DatabendUnaryPostfixOperation;
3445
import sqlancer.databend.ast.DatabendUnaryPostfixOperation.DatabendUnaryPostfixOperator;
3546
import sqlancer.databend.ast.DatabendUnaryPrefixOperation;
3647
import sqlancer.databend.ast.DatabendUnaryPrefixOperation.DatabendUnaryPrefixOperator;
3748

3849
public class DatabendNewExpressionGenerator
39-
extends TypedExpressionGenerator<DatabendExpression, DatabendColumn, DatabendDataType> {
50+
extends TypedExpressionGenerator<DatabendExpression, DatabendColumn, DatabendDataType>
51+
implements NoRECGenerator<DatabendSelect, DatabendJoin, DatabendExpression, DatabendTable, DatabendColumn> {
4052

4153
private final DatabendGlobalState globalState;
54+
private List<DatabendTable> tables;
4255

4356
private final int maxDepth;
4457
private boolean allowAggregateFunctions;
@@ -248,7 +261,8 @@ DatabendExpression getLike(int depth, DatabendDataType dataType) {
248261
}
249262

250263
public DatabendExpression generateExpressionWithExpectedResult(DatabendDataType type) {
251-
// DatabendNewExpressionGenerator gen = new DatabendNewExpressionGenerator(globalState).setColumns(columns);
264+
// DatabendNewExpressionGenerator gen = new
265+
// DatabendNewExpressionGenerator(globalState).setColumns(columns);
252266
// gen.setRowValue(rowValue);
253267
DatabendExpression expr;
254268
do {
@@ -340,4 +354,71 @@ public DatabendExpression generateHavingClause() {
340354
return expression;
341355
}
342356

357+
@Override
358+
public NoRECGenerator<DatabendSelect, DatabendJoin, DatabendExpression, DatabendTable, DatabendColumn> setTablesAndColumns(
359+
AbstractTables<DatabendTable, DatabendColumn> tables) {
360+
this.columns = tables.getColumns();
361+
this.tables = tables.getTables();
362+
363+
return this;
364+
}
365+
366+
@Override
367+
public DatabendExpression generateBooleanExpression() {
368+
return generateExpression(DatabendDataType.BOOLEAN);
369+
}
370+
371+
@Override
372+
public DatabendSelect generateSelect() {
373+
return new DatabendSelect();
374+
}
375+
376+
@Override
377+
public List<DatabendJoin> getRandomJoinClauses() {
378+
List<DatabendTableReference> tableList = tables.stream().map(t -> new DatabendTableReference(t))
379+
.collect(Collectors.toList());
380+
List<DatabendJoin> joins = DatabendJoin.getJoins(tableList, globalState);
381+
tables = tableList.stream().map(t -> t.getTable()).collect(Collectors.toList());
382+
return joins;
383+
}
384+
385+
@Override
386+
public List<DatabendExpression> getTableRefs() {
387+
return tables.stream().map(t -> new DatabendTableReference(t)).collect(Collectors.toList());
388+
}
389+
390+
@Override
391+
public String generateOptimizedQueryString(DatabendSelect select, DatabendExpression whereCondition,
392+
boolean shouldUseAggregate) {
393+
if (shouldUseAggregate) {
394+
DatabendExpression aggr = new DatabendAggregateOperation(
395+
List.of(new DatabendColumnReference(new DatabendColumn("*",
396+
new DatabendCompositeDataType(DatabendDataType.INT, 0), false, false))),
397+
DatabendAggregateFunction.COUNT);
398+
select.setFetchColumns(List.of(aggr));
399+
} else {
400+
List<DatabendExpression> allColumns = columns.stream().map((c) -> new DatabendColumnReference(c))
401+
.collect(Collectors.toList());
402+
select.setFetchColumns(allColumns);
403+
if (Randomly.getBooleanWithSmallProbability()) {
404+
select.setOrderByClauses(generateOrderBys());
405+
}
406+
}
407+
select.setWhereClause(whereCondition);
408+
409+
return select.asString();
410+
}
411+
412+
@Override
413+
public String generateUnoptimizedQueryString(DatabendSelect select, DatabendExpression whereCondition) {
414+
DatabendExpression asText = new DatabendPostFixText(new DatabendCastOperation(
415+
new DatabendPostFixText(whereCondition,
416+
" IS NOT NULL AND " + DatabendToStringVisitor.asString(whereCondition)),
417+
new DatabendCompositeDataType(DatabendDataType.INT, 8)), "as count");
418+
select.setFetchColumns(List.of(asText));
419+
select.setWhereClause(null);
420+
421+
return "SELECT SUM(count) FROM (" + select.asString() + ") as res";
422+
}
423+
343424
}

‎src/sqlancer/databend/gen/DatabendRandomQuerySynthesizer.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,8 @@ public static DatabendSelect generateSelect(DatabendGlobalState globalState, int
5151
List<DatabendTable> tables = targetTables.getTables();
5252
List<DatabendTableReference> tableList = tables.stream().map(t -> new DatabendTableReference(t))
5353
.collect(Collectors.toList());
54-
List<DatabendExpression> joins = DatabendJoin.getJoins(tableList, globalState);
55-
select.setJoinList(joins);
54+
List<DatabendJoin> joins = DatabendJoin.getJoins(tableList, globalState);
55+
select.setJoinList(joins.stream().collect(Collectors.toList()));
5656
select.setFromList(tableList.stream().collect(Collectors.toList()));
5757
if (Randomly.getBoolean()) {
5858
select.setWhereClause(gen.generateExpression(DatabendSchema.DatabendDataType.BOOLEAN));

‎src/sqlancer/databend/test/tlp/DatabendQueryPartitioningBase.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ public void check() throws SQLException {
6262
List<DatabendTableReference> tableList = tables.stream().map(t -> new DatabendTableReference(t))
6363
.collect(Collectors.toList());
6464
if (!DatabendBugs.bug9236) {
65-
List<DatabendExpression> joins = DatabendJoin.getJoins(tableList, state);
65+
List<DatabendJoin> joins = DatabendJoin.getJoins(tableList, state);
6666
select.setJoinList(joins.stream().collect(Collectors.toList()));
6767
}
6868
select.setFromList(tableList.stream().collect(Collectors.toList()));

0 commit comments

Comments
 (0)