Skip to content

Commit cc152b5

Browse files
yuancuyuancu
committed
Refactor: pass expr values to udf implementations instead of passing java object and expr type separately
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 9147dfe commit cc152b5

13 files changed

Lines changed: 230 additions & 372 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/utils/UserDefinedFunctionUtils.java‎

Lines changed: 14 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -10,17 +10,11 @@
1010
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT.*;
1111

1212
import java.time.Instant;
13-
import java.time.LocalDateTime;
1413
import java.time.ZoneId;
15-
import java.time.format.DateTimeFormatter;
1614
import java.util.ArrayList;
17-
import java.util.Arrays;
1815
import java.util.Collections;
1916
import java.util.List;
20-
import java.util.Locale;
21-
import java.util.Objects;
2217
import java.util.TimeZone;
23-
import java.util.stream.Collectors;
2418
import org.apache.calcite.DataContext;
2519
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
2620
import org.apache.calcite.adapter.enumerable.NullPolicy;
@@ -94,29 +88,6 @@ static SqlReturnTypeInference getReturnTypeInferenceForArray() {
9488
};
9589
}
9690

97-
/** Check whether the given array contains null values. */
98-
public static boolean containsNull(Object[] objects) {
99-
return Arrays.stream(objects).anyMatch(Objects::isNull);
100-
}
101-
102-
public static String formatTimestampWithoutUnnecessaryNanos(LocalDateTime localDateTime) {
103-
String base = localDateTime.format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"));
104-
int nano = localDateTime.getNano();
105-
if (nano == 0) return base;
106-
107-
String nanoStr = String.format(Locale.ENGLISH, "%09d", nano);
108-
nanoStr = nanoStr.replaceFirst("0+$", "");
109-
if (!nanoStr.isEmpty()) {
110-
return base + "." + nanoStr;
111-
}
112-
return base;
113-
}
114-
115-
public static ExprType transferDateRelatedTimeName(RexNode candidate) {
116-
RelDataType type = candidate.getType();
117-
return OpenSearchTypeFactory.convertRelDataTypeToExprType(type);
118-
}
119-
12091
// TODO: pass the function properties directly to the UDF instead of string
12192
public static FunctionProperties restoreFunctionProperties(DataContext dataContext) {
12293
long currentTimeInNanos = DataContext.Variable.UTC_TIMESTAMP.get(dataContext);
@@ -128,16 +99,18 @@ public static FunctionProperties restoreFunctionProperties(DataContext dataConte
12899
return new FunctionProperties(instant, zoneId, QueryType.PPL);
129100
}
130101

131-
public static List<Expression> addTypeAndContext(
132-
List<Expression> operands, RexCall rexCall, Expression root) {
133-
// Box primitive values so that they can match signatures with boxed types
134-
List<Expression> enrichedOperands =
135-
operands.stream().map(Expressions::box).collect(Collectors.toList());
136-
for (RexNode rexNode : rexCall.getOperands()) {
137-
enrichedOperands.add(Expressions.constant(transferDateRelatedTimeName(rexNode)));
138-
}
139-
enrichedOperands.add(root);
140-
return enrichedOperands;
102+
/**
103+
* Convert java objects to ExprValue, so that the parameters fit the expr function signature. It
104+
* invokes ExprValueUtils.fromObjectValue to convert the java objects to ExprValue. Note that
105+
* date/time/timestamp strings will be converted to strings instead of ExprDateValue, etc.
106+
*
107+
* @param operands the operands to convert
108+
* @param rexCall the RexCall object containing the operands
109+
* @return the converted operands
110+
*/
111+
public static List<Expression> convertToExprValues(List<Expression> operands, RexCall rexCall) {
112+
List<RelDataType> types = rexCall.getOperands().stream().map(RexNode::getType).toList();
113+
return convertToExprValues(operands, types);
141114
}
142115

143116
/**
@@ -198,7 +171,7 @@ public SqlReturnTypeInference getReturnTypeInference() {
198171
};
199172
}
200173

201-
private static List<Expression> prependTimestampAsProperty(
174+
public static List<Expression> prependFunctionProperties(
202175
List<Expression> operands, RexToLixTranslator translator) {
203176
List<Expression> operandsWithProperties = new ArrayList<>(operands);
204177
Expression properties =
@@ -218,8 +191,7 @@ public static ImplementorUDF adaptExprMethodWithPropertiesToUDF(
218191
List<Expression> operands =
219192
convertToExprValues(
220193
translatedOperands, call.getOperands().stream().map(RexNode::getType).toList());
221-
List<Expression> operandsWithProperties =
222-
prependTimestampAsProperty(operands, translator);
194+
List<Expression> operandsWithProperties = prependFunctionProperties(operands, translator);
223195
Expression exprResult = Expressions.call(type, methodName, operandsWithProperties);
224196
return Expressions.call(exprResult, "valueForCalcite");
225197
};

‎core/src/main/java/org/opensearch/sql/calcite/utils/datetime/DateTimeApplyUtils.java‎

Lines changed: 0 additions & 76 deletions
This file was deleted.

‎core/src/main/java/org/opensearch/sql/calcite/utils/datetime/DateTimeConversionUtils.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.utils.datetime;
7+
8+
import java.time.Duration;
9+
import java.time.Period;
10+
import java.time.temporal.TemporalAmount;
11+
import org.apache.calcite.avatica.util.TimeUnit;
12+
import org.opensearch.sql.data.model.*;
13+
import org.opensearch.sql.data.type.ExprCoreType;
14+
import org.opensearch.sql.exception.SemanticCheckException;
15+
import org.opensearch.sql.expression.function.FunctionProperties;
16+
17+
public final class DateTimeConversionUtils {
18+
private DateTimeConversionUtils() {}
19+
20+
/**
21+
* Convert the given ExprValue to an ExprTimestampValue. If the input is a string, it will convert
22+
* date / time / timestamp strings to ExprTimestampValue.
23+
*
24+
* @param value the value to convert, can be either a ExprDateValue, ExprTimeValue,
25+
* ExprTimestampValue or ExprStringValue
26+
* @param properties the function properties
27+
* @return the converted ExprTimestampValue
28+
*/
29+
public static ExprTimestampValue forceConvertToTimestampValue(
30+
ExprValue value, FunctionProperties properties) {
31+
return switch (value) {
32+
case ExprTimestampValue timestampValue -> timestampValue;
33+
case ExprDateValue dateValue -> (ExprTimestampValue)
34+
ExprValueUtils.timestampValue(dateValue.timestampValue());
35+
case ExprTimeValue timeValue -> (ExprTimestampValue)
36+
ExprValueUtils.timestampValue(timeValue.timestampValue(properties));
37+
case ExprStringValue stringValue -> new ExprTimestampValue(
38+
DateTimeParser.parse(stringValue.stringValue()));
39+
default -> throw new SemanticCheckException(
40+
String.format(
41+
"Cannot convert %s to timestamp, only STRING, DATE, TIME and TIMESTAMP are supported",
42+
value.type()));
43+
};
44+
}
45+
46+
/**
47+
* Convert the given ExprValue to an ExprTimestampValue. If the input is a string, it only accepts
48+
* a string formatted as a valid timestamp 'yyyy-MM-dd HH:mm:ss[.SSSSSSSSS]'.
49+
*
50+
* @param value the value to convert, can be either a ExprDateValue, ExprTimeValue,
51+
* ExprTimestampValue or ExprStringValue
52+
* @param properties the function properties
53+
* @return the converted ExprTimestampValue
54+
*/
55+
public static ExprTimestampValue convertToTimestampValue(
56+
ExprValue value, FunctionProperties properties) {
57+
if (value instanceof ExprTimestampValue timestampValue) {
58+
return timestampValue;
59+
} else if (value instanceof ExprTimeValue) {
60+
return new ExprTimestampValue(((ExprTimeValue) value).timestampValue(properties));
61+
} else if (value.type() == ExprCoreType.STRING) {
62+
return new ExprTimestampValue(value.stringValue());
63+
} else {
64+
try {
65+
return new ExprTimestampValue(value.timestampValue());
66+
} catch (SemanticCheckException e) {
67+
throw new SemanticCheckException(
68+
String.format(
69+
"Cannot convert %s to timestamp, only STRING, DATE, TIME and TIMESTAMP are"
70+
+ " supported",
71+
value.type()),
72+
e);
73+
}
74+
}
75+
}
76+
77+
/**
78+
* Convert the given ExprValue to an ExprDateValue. If the input is a string, it only accepts a
79+
* string formatted as a valid date 'yyyy-MM-dd'.
80+
*
81+
* @param value the value to convert, can be either a ExprDateValue, ExprTimeValue,
82+
* ExprTimestampValue or ExprStringValue
83+
* @param properties the function properties
84+
* @return the converted ExprDateValue
85+
*/
86+
public static ExprDateValue convertToDateValue(ExprValue value, FunctionProperties properties) {
87+
switch (value) {
88+
case ExprDateValue dateValue -> {
89+
return dateValue;
90+
}
91+
case ExprTimeValue timeValue -> {
92+
return new ExprDateValue(timeValue.dateValue(properties));
93+
}
94+
case ExprTimestampValue timestampValue -> {
95+
return new ExprDateValue(timestampValue.dateValue());
96+
}
97+
case ExprStringValue ignored -> {
98+
return new ExprDateValue(value.stringValue());
99+
}
100+
default -> {
101+
throw new SemanticCheckException(
102+
String.format(
103+
"Cannot convert %s to date, only STRING, DATE, TIME and TIMESTAMP are supported",
104+
value.type()));
105+
}
106+
}
107+
}
108+
109+
/**
110+
* Create a temporal amount of the given number of units. For duration below a day, it returns
111+
* duration; for duration including and above a day, it returns period for natural days, months,
112+
* quarters, and years, which may be of unfixed lengths.
113+
*
114+
* @param number The count of unit
115+
* @param unit The unit of the temporal amount
116+
* @return A temporal amount value, can be either a Period or a Duration
117+
*/
118+
public static TemporalAmount convertToTemporalAmount(long number, TimeUnit unit) {
119+
return switch (unit) {
120+
case YEAR -> Period.ofYears((int) number);
121+
case QUARTER -> Period.ofMonths((int) number * 3);
122+
case MONTH -> Period.ofMonths((int) number);
123+
case WEEK -> Period.ofWeeks((int) number);
124+
case DAY -> Period.ofDays((int) number);
125+
case HOUR -> Duration.ofHours(number);
126+
case MINUTE -> Duration.ofMinutes(number);
127+
case SECOND -> Duration.ofSeconds(number);
128+
case MILLISECOND -> Duration.ofMillis(number);
129+
case MICROSECOND -> Duration.ofNanos(number * 1000);
130+
case NANOSECOND -> Duration.ofNanos(number);
131+
132+
default -> throw new UnsupportedOperationException(
133+
"No mapping defined for Calcite TimeUnit: " + unit);
134+
};
135+
}
136+
}

0 commit comments

Comments
 (0)