Skip to content

Commit f3f94ce

Browse files
committed
Add ITs for UDF type checking
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent fe0385c commit f3f94ce

7 files changed

Lines changed: 171 additions & 33 deletions

File tree

core/src/main/java/org/opensearch/sql/exception/CalciteUnsupportedException.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,8 @@ public class CalciteUnsupportedException extends QueryEngineException {
1010
public CalciteUnsupportedException(String message) {
1111
super(message);
1212
}
13+
14+
public CalciteUnsupportedException(String message, Throwable cause) {
15+
super(message, cause);
16+
}
1317
}

core/src/main/java/org/opensearch/sql/executor/QueryService.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ public void executeWithCalcite(
108108
} else {
109109
if (t instanceof Error) {
110110
// Calcite may throw AssertError during query execution.
111-
listener.onFailure(new CalciteUnsupportedException(t.getMessage()));
111+
listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
112112
} else {
113113
listener.onFailure((Exception) t);
114114
}

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -242,7 +242,6 @@ void populate() {
242242
registerOperator(ADD, SqlStdOperatorTable.PLUS);
243243
registerOperator(SUBTRACT, SqlStdOperatorTable.MINUS);
244244
registerOperator(MULTIPLY, SqlStdOperatorTable.MULTIPLY);
245-
// registerOperator(DIVIDE, SqlStdOperatorTable.DIVIDE);
246245
registerOperator(ASCII, SqlStdOperatorTable.ASCII);
247246
registerOperator(LENGTH, SqlStdOperatorTable.CHAR_LENGTH);
248247
registerOperator(LOWER, SqlStdOperatorTable.LOWER);

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLBuiltinFunctionIT.java

Lines changed: 60 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -56,19 +56,6 @@ public void testSqrtNegativeArgShouldReturnNull() {
5656
verifyDataRows(actual, rows((Object) null));
5757
}
5858

59-
@Test
60-
public void testSqrtStrArgShouldThrowException() {
61-
Exception nanException =
62-
assertThrows(
63-
ExpressionEvaluationException.class,
64-
() ->
65-
executeQuery(
66-
String.format(
67-
"source=%s | head 1 | eval sqrt_name = sqrt(name) | fields sqrt_name",
68-
TEST_INDEX_STATE_COUNTRY)));
69-
verifyErrorMessageContains(nanException, "SQRT function expects {[NUMERIC]}, but got [STRING]");
70-
}
71-
7259
@Test
7360
public void testSinAndCosAndAsinAndAcos() {
7461
JSONObject actual =
@@ -317,22 +304,6 @@ public void testModByZeroShouldReturnNull() {
317304
verifyDataRows(actual, rows((Object) null));
318305
}
319306

320-
@Test
321-
public void testMod3ArgsShouldThrowExprEvalException() {
322-
Exception wrongArgException =
323-
assertThrows(
324-
ExpressionEvaluationException.class,
325-
() ->
326-
executeQuery(
327-
String.format(
328-
"source=%s | eval z = mod(float_number, integer_number, byte_number) |"
329-
+ " fields z",
330-
TEST_INDEX_DATATYPE_NUMERIC)));
331-
verifyErrorMessageContains(
332-
wrongArgException,
333-
"MOD function expects {[NUMERIC,NUMERIC]}, but got [FLOAT,INTEGER,BYTE]");
334-
}
335-
336307
@Test
337308
public void testRadiansAndDegrees() {
338309
JSONObject actual =
@@ -440,4 +411,64 @@ public void testDivideShouldReturnNull() {
440411
schema("r6", "double"));
441412
verifyDataRows(actual, rows(null, null, null, null, null));
442413
}
414+
415+
// Test type checking on udf with direct registration: register(SQRT, funcImp)
416+
@Test
417+
public void testSqrtWrongArgShouldThrow() {
418+
Exception nanException =
419+
assertThrows(
420+
ExpressionEvaluationException.class,
421+
() ->
422+
executeQuery(
423+
String.format(
424+
"source=%s | head 1 | eval sqrt_name = sqrt(name) | fields sqrt_name",
425+
TEST_INDEX_STATE_COUNTRY)));
426+
verifyErrorMessageContains(nanException, "SQRT function expects {[NUMERIC]}, but got [STRING]");
427+
}
428+
429+
// Test UDF registered with PPL builtin operators: registerOperator(MOD, PPLBuiltinOperators.MOD);
430+
@Test
431+
public void testModWrongArgShouldThrow() {
432+
Exception wrongArgException =
433+
assertThrows(
434+
ExpressionEvaluationException.class,
435+
() ->
436+
executeQuery(
437+
String.format(
438+
"source=%s | eval z = mod(float_number, integer_number, byte_number) |"
439+
+ " fields z",
440+
TEST_INDEX_DATATYPE_NUMERIC)));
441+
verifyErrorMessageContains(
442+
wrongArgException,
443+
"MOD function expects {[NUMERIC,NUMERIC]}, but got [FLOAT,INTEGER,BYTE]");
444+
}
445+
446+
// Test UDF registered with sql std operators: registerOperator(PI, SqlStdOperatorTable.PI)
447+
@Test
448+
public void testPiWithArgShouldThrow() {
449+
Exception wrongArgException =
450+
assertThrows(
451+
ExpressionEvaluationException.class,
452+
() ->
453+
executeQuery(
454+
String.format(
455+
"source=%s | eval pi = pi(1) | fields pi", TEST_INDEX_STATE_COUNTRY)));
456+
verifyErrorMessageContains(wrongArgException, "PI function expects {[]}, but got [INTEGER]");
457+
}
458+
459+
// Test UDF registered with sql library operators: registerOperator(LOG2,
460+
// SqlLibraryOperators.LOG2)
461+
@Test
462+
public void testLog2WithWrongArgShouldThrow() {
463+
Exception wrongArgException =
464+
assertThrows(
465+
ExpressionEvaluationException.class,
466+
() ->
467+
executeQuery(
468+
String.format(
469+
"source=%s | eval log2 = log2(age, 2) | fields log2",
470+
TEST_INDEX_STATE_COUNTRY)));
471+
verifyErrorMessageContains(
472+
wrongArgException, "LOG2 function expects {[NUMERIC]}, but got [INTEGER,INTEGER]");
473+
}
443474
}

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLCryptographicFunctionIT.java

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.io.IOException;
1616
import org.json.JSONObject;
1717
import org.junit.jupiter.api.Test;
18+
import org.opensearch.sql.exception.ExpressionEvaluationException;
1819

1920
public class CalcitePPLCryptographicFunctionIT extends CalcitePPLIntegTestCase {
2021
@Override
@@ -88,4 +89,18 @@ public void testSha2WrongAlgorithmShouldThrow() {
8889
TEST_INDEX_STATE_COUNTRY)));
8990
verifyErrorMessageContains(e, "Unsupported SHA2 algorithm");
9091
}
92+
93+
@Test
94+
public void testSha2WrongArgShouldThrow() {
95+
Throwable e =
96+
assertThrows(
97+
ExpressionEvaluationException.class,
98+
() ->
99+
executeQuery(
100+
String.format(
101+
"source=%s | head 1 | eval sha256 = SHA2('hello', '256') | fields sha256",
102+
TEST_INDEX_STATE_COUNTRY)));
103+
verifyErrorMessageContains(
104+
e, "SHA2 function expects {[STRING,INTEGER]}, but got [STRING,STRING]");
105+
}
91106
}

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLDateTimeBuiltinFunctionIT.java

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.opensearch.client.Request;
3333
import org.opensearch.sql.data.model.ExprDateValue;
3434
import org.opensearch.sql.data.model.ExprIntegerValue;
35+
import org.opensearch.sql.exception.ExpressionEvaluationException;
3536
import org.opensearch.sql.expression.function.FunctionProperties;
3637

3738
public class CalcitePPLDateTimeBuiltinFunctionIT extends CalcitePPLIntegTestCase {
@@ -100,7 +101,6 @@ public void testDate() {
100101
Date.valueOf("1984-04-12")));
101102
}
102103

103-
// TODO: Add Hook to fix the input timestamp. We should actually add HOOK in init.
104104
@Test
105105
public void testTimestampWithTimeInput() {
106106
String utcTomorrow = LocalDate.now().plusDays(1).toString();
@@ -169,6 +169,23 @@ public void testTimestamp() {
169169
"2009-12-12 13:40:04.123456789"));
170170
}
171171

172+
@Test
173+
public void testTimestampWithWrongArgShouldThrow() {
174+
Exception e =
175+
assertThrows(
176+
ExpressionEvaluationException.class,
177+
() ->
178+
executeQuery(
179+
String.format(
180+
"source=%s | eval timestamp = timestamp('2020-08-26 13:49:00', 2009)"
181+
+ "| fields timestamp | head 1",
182+
TEST_INDEX_DATE_FORMATS)));
183+
verifyErrorMessageContains(
184+
e,
185+
"TIMESTAMP function expects {[STRING], [DATETIME], [STRING,STRING], [DATETIME,DATETIME],"
186+
+ " [STRING,DATETIME], [DATETIME,STRING]}, but got [STRING,INTEGER]");
187+
}
188+
172189
@Test
173190
public void testTime() {
174191
JSONObject actual =
@@ -777,6 +794,19 @@ public void testCurrentDateTimeWithComparison() {
777794
verifyNumOfRows(actual, 7);
778795
}
779796

797+
@Test
798+
public void testCurDateWithArgShouldThrow() {
799+
Exception e =
800+
assertThrows(
801+
ExpressionEvaluationException.class,
802+
() ->
803+
executeQuery(
804+
String.format(
805+
"source=%s | eval curdate = CURDATE(1) | fields curdate | head 1",
806+
TEST_INDEX_DATE_FORMATS)));
807+
verifyErrorMessageContains(e, "CURDATE function expects {[]}, but got [INTEGER]");
808+
}
809+
780810
@Test
781811
public void testUtc() {
782812
JSONObject actual =

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLStringBuiltinFunctionIT.java

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.json.JSONObject;
1515
import org.junit.jupiter.api.Test;
1616
import org.opensearch.client.Request;
17+
import org.opensearch.sql.exception.ExpressionEvaluationException;
1718

1819
public class CalcitePPLStringBuiltinFunctionIT extends CalcitePPLIntegTestCase {
1920
@Override
@@ -45,7 +46,7 @@ public void testConcat() {
4546
JSONObject actual =
4647
executeQuery(
4748
String.format(
48-
"source=%s | where name=concat('He', 'llo') | fields name, age",
49+
"source=%s | where name=concat('He', 'll', 'o') | fields name, age",
4950
TEST_INDEX_STATE_COUNTRY));
5051

5152
verifySchema(actual, schema("name", "string"), schema("age", "integer"));
@@ -126,6 +127,20 @@ public void testLower() {
126127
verifyDataRows(actual, rows("Hello", 30));
127128
}
128129

130+
// Test registered Sql Std Operator: registerOperator(funcName, SqlStdOperatorTable.OPERATOR)
131+
@Test
132+
public void testLowerWrongArgShouldThrow() {
133+
Exception e =
134+
assertThrows(
135+
ExpressionEvaluationException.class,
136+
() ->
137+
executeQuery(
138+
String.format(
139+
"source=%s | where lower(age) = 'hello' | fields name, age",
140+
TEST_INDEX_STATE_COUNTRY)));
141+
verifyErrorMessageContains(e, "LOWER function expects {[CHARACTER]}, but got [INTEGER]");
142+
}
143+
129144
@Test
130145
public void testUpper() {
131146
JSONObject actual =
@@ -248,6 +263,20 @@ public void testLTrim() throws IOException {
248263
verifyDataRows(actual, rows(" Jim", 27));
249264
}
250265

266+
// Test directly registered UDF: register(funcname, FuncImp)
267+
@Test
268+
public void testLtrimWrongArgShouldThrow() {
269+
Exception e =
270+
assertThrows(
271+
ExpressionEvaluationException.class,
272+
() ->
273+
executeQuery(
274+
String.format(
275+
"source=%s | where ltrim(year, age) = 'Jim' | fields name, age",
276+
TEST_INDEX_STATE_COUNTRY)));
277+
verifyErrorMessageContains(e, "LTRIM function expects {[STRING]}, but got [INTEGER,INTEGER]");
278+
}
279+
251280
@Test
252281
public void testReverse() throws IOException {
253282
Request request1 =
@@ -266,6 +295,21 @@ public void testReverse() throws IOException {
266295
verifyDataRows(actual, rows("DeD", 27));
267296
}
268297

298+
// Test udf registered via sql library operator: registerOperator(REVERSE,
299+
// SqlLibraryOperators.REVERSE);
300+
@Test
301+
public void testReverseWrongArgShouldThrow() {
302+
Exception e =
303+
assertThrows(
304+
ExpressionEvaluationException.class,
305+
() ->
306+
executeQuery(
307+
String.format(
308+
"source=%s | where reverse(year) = '3202' | fields year",
309+
TEST_INDEX_STATE_COUNTRY)));
310+
verifyErrorMessageContains(e, "REVERSE function expects {[CHARACTER]}, but got [INTEGER]");
311+
}
312+
269313
@Test
270314
public void testRight() throws IOException {
271315
Request request1 =
@@ -325,6 +369,21 @@ public void testStrCmp() throws IOException {
325369
verifyDataRows(actual, rows("Jane", 20));
326370
}
327371

372+
// test type checking on UDF with direct registration: register(funcname, FuncImp)
373+
@Test
374+
public void testStrCmpWrongArgShouldThrow() {
375+
Exception e =
376+
assertThrows(
377+
ExpressionEvaluationException.class,
378+
() ->
379+
executeQuery(
380+
String.format(
381+
"source=%s | where strcmp(age, 'Jane') = 0 | fields name, age",
382+
TEST_INDEX_STATE_COUNTRY)));
383+
verifyErrorMessageContains(
384+
e, "STRCMP function expects {[STRING,STRING]}, but got [INTEGER,STRING]");
385+
}
386+
328387
private void prepareTrim() throws IOException {
329388
Request request1 =
330389
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/5?refresh=true");

0 commit comments

Comments
 (0)