-
Notifications
You must be signed in to change notification settings - Fork 397
Expand file tree
/
Copy pathPrestoFunction.java
More file actions
126 lines (109 loc) · 5.16 KB
/
PrestoFunction.java
File metadata and controls
126 lines (109 loc) · 5.16 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
package sqlancer.presto.ast;
import java.util.ArrayList;
import java.util.List;
import sqlancer.Randomly;
import sqlancer.presto.PrestoSchema;
import sqlancer.presto.gen.PrestoTypedExpressionGenerator;
public interface PrestoFunction extends PrestoExpression {
String getFunctionName();
boolean isCompatibleWithReturnType(PrestoSchema.PrestoCompositeDataType returnType);
PrestoSchema.PrestoDataType[] getArgumentTypes(PrestoSchema.PrestoCompositeDataType returnType);
default List<PrestoExpression> getArgumentsForReturnType(PrestoTypedExpressionGenerator gen, int depth,
PrestoSchema.PrestoDataType[] argumentTypes, PrestoSchema.PrestoCompositeDataType returnType) {
List<PrestoExpression> arguments = new ArrayList<>();
// This is a workaround based on the assumption that array types should refer to
// the same element type.
PrestoSchema.PrestoCompositeDataType savedArrayType = null;
if (returnType.getPrimitiveDataType() == PrestoSchema.PrestoDataType.ARRAY) {
savedArrayType = returnType;
}
// -1 - unlimited number of arguments
if (getNumberOfArguments() == -1) {
PrestoSchema.PrestoDataType dataType = argumentTypes[0];
// TODO: consider upper
long no = Randomly.getNotCachedInteger(2, 10);
for (int i = 0; i < no; i++) {
PrestoSchema.PrestoCompositeDataType type;
if (dataType == PrestoSchema.PrestoDataType.ARRAY) {
if (savedArrayType == null) {
savedArrayType = dataType.get();
}
type = savedArrayType;
} else {
type = PrestoSchema.PrestoCompositeDataType.fromDataType(dataType);
}
arguments.add(gen.generateExpression(type, depth + 1));
}
} else {
for (PrestoSchema.PrestoDataType arg : argumentTypes) {
PrestoSchema.PrestoCompositeDataType dataType;
if (arg == PrestoSchema.PrestoDataType.ARRAY) {
if (savedArrayType == null) {
savedArrayType = arg.get();
}
dataType = savedArrayType;
} else {
dataType = PrestoSchema.PrestoCompositeDataType.fromDataType(arg);
}
PrestoExpression expression = gen.generateExpression(dataType, depth + 1);
arguments.add(expression);
}
}
return arguments;
}
default List<PrestoExpression> getArgumentsForReturnType(PrestoTypedExpressionGenerator gen, int depth,
PrestoSchema.PrestoCompositeDataType returnType, boolean orderable) {
List<PrestoExpression> arguments = new ArrayList<>();
// This is a workaround based on the assumption that array types should refer to
// the same element type.
PrestoSchema.PrestoCompositeDataType savedArrayType = null;
if (returnType.getPrimitiveDataType() == PrestoSchema.PrestoDataType.ARRAY) {
savedArrayType = returnType;
}
if (getNumberOfArguments() == -1) {
PrestoSchema.PrestoDataType dataType = getArgumentTypes(returnType)[0];
// TODO: consider upper
long no = Randomly.getNotCachedInteger(2, 10);
for (int i = 0; i < no; i++) {
PrestoSchema.PrestoCompositeDataType compositeDataType;
if (dataType == PrestoSchema.PrestoDataType.ARRAY) {
if (savedArrayType == null) {
savedArrayType = dataType.get();
}
compositeDataType = savedArrayType;
} else {
compositeDataType = PrestoSchema.PrestoCompositeDataType.fromDataType(dataType);
}
arguments.add(gen.generateExpression(compositeDataType, depth + 1));
}
} else {
for (PrestoSchema.PrestoDataType dataType : getArgumentTypes(returnType)) {
PrestoSchema.PrestoCompositeDataType compositeDataType;
if (dataType == PrestoSchema.PrestoDataType.ARRAY) {
if (savedArrayType == null) {
PrestoSchema.PrestoCompositeDataType arrayType;
do {
arrayType = dataType.get();
} while (!arrayType.getElementType().isOrderable());
savedArrayType = arrayType;
}
compositeDataType = savedArrayType;
} else {
compositeDataType = PrestoSchema.PrestoCompositeDataType.fromDataType(dataType);
}
PrestoExpression expression = gen.generateExpression(compositeDataType, depth + 1);
arguments.add(expression);
}
}
return arguments;
}
default int getNumberOfArguments() {
return getArgumentTypes(null).length;
}
default boolean shouldPreserveOrderOfArguments() {
return false;
}
default boolean isStandardFunction() {
return true;
}
}