forked from opendistro-for-elasticsearch/sql
-
Notifications
You must be signed in to change notification settings - Fork 0
/
CaseClause.java
117 lines (101 loc) · 4.02 KB
/
CaseClause.java
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
/*
* Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* or in the "license" file accompanying this file. This file is distributed
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing
* permissions and limitations under the License.
*
*/
package com.amazon.opendistroforelasticsearch.sql.expression.conditional.cases;
import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.UNDEFINED;
import static com.amazon.opendistroforelasticsearch.sql.data.type.ExprCoreType.UNKNOWN;
import com.amazon.opendistroforelasticsearch.sql.data.model.ExprNullValue;
import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue;
import com.amazon.opendistroforelasticsearch.sql.data.type.ExprType;
import com.amazon.opendistroforelasticsearch.sql.expression.Expression;
import com.amazon.opendistroforelasticsearch.sql.expression.ExpressionNodeVisitor;
import com.amazon.opendistroforelasticsearch.sql.expression.FunctionExpression;
import com.amazon.opendistroforelasticsearch.sql.expression.env.Environment;
import com.amazon.opendistroforelasticsearch.sql.expression.function.FunctionName;
import com.google.common.collect.ImmutableList;
import java.util.List;
import java.util.stream.Collectors;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.ToString;
/**
* A CASE clause is very different from a regular function. Functions have well-defined signature,
* though CASE clause is more like a function implementation which requires type check "manually".
*/
@EqualsAndHashCode(callSuper = false)
@Getter
@ToString
public class CaseClause extends FunctionExpression {
/**
* List of WHEN clauses.
*/
private final List<WhenClause> whenClauses;
/**
* Default result if none of WHEN conditions match.
*/
private final Expression defaultResult;
/**
* Initialize case clause.
*/
public CaseClause(List<WhenClause> whenClauses, Expression defaultResult) {
super(FunctionName.of("case"), concatArgs(whenClauses, defaultResult));
this.whenClauses = whenClauses;
this.defaultResult = defaultResult;
}
@Override
public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
for (WhenClause when : whenClauses) {
if (when.isTrue(valueEnv)) {
return when.valueOf(valueEnv);
}
}
return (defaultResult == null) ? ExprNullValue.of() : defaultResult.valueOf(valueEnv);
}
@Override
public ExprType type() {
List<ExprType> types = allResultTypes();
// Return undefined if all WHEN/ELSE return NULL
return types.isEmpty() ? UNDEFINED : types.get(0);
}
@Override
public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
return visitor.visitCase(this, context);
}
/**
* Get types of each result in WHEN clause and ELSE clause.
* Exclude UNKNOWN type from NULL literal which means NULL in THEN or ELSE clause
* is not included in result.
* @return all result types. Use list so caller can generate friendly error message.
*/
public List<ExprType> allResultTypes() {
List<ExprType> types = whenClauses.stream()
.map(WhenClause::type)
.collect(Collectors.toList());
if (defaultResult != null) {
types.add(defaultResult.type());
}
types.removeIf(type -> (type == UNDEFINED));
return types;
}
private static List<Expression> concatArgs(List<WhenClause> whenClauses,
Expression defaultResult) {
ImmutableList.Builder<Expression> args = ImmutableList.builder();
whenClauses.forEach(args::add);
if (defaultResult != null) {
args.add(defaultResult);
}
return args.build();
}
}