Skip to content

Commit

Permalink
Add string builtin functions (#3393)
Browse files Browse the repository at this point in the history 
* add string udfs

Signed-off-by: xinyual <[email protected]>

* add it to string

Signed-off-by: xinyual <[email protected]>

* add IT for string function

Signed-off-by: xinyual <[email protected]>

* remove change for local test

Signed-off-by: xinyual <[email protected]>

* revert change

Signed-off-by: xinyual <[email protected]>

---------

Signed-off-by: xinyual <[email protected]>
  • Loading branch information
@xinyual
xinyual authored Mar 10, 2025
1 parent dcf2057 commit b7f6cf7
Show file tree
Hide file tree
Showing 2 changed files with 284 additions and 2 deletions.
47 changes: 45 additions & 2 deletions core/src/main/java/org/opensearch/sql/calcite/utils/BuiltinFunctionUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import org.apache.calcite.sql.SqlOperator;
import org.apache.calcite.sql.fun.SqlLibraryOperators;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.fun.SqlTrimFunction;
import org.apache.calcite.sql.type.ReturnTypes;
import org.opensearch.sql.calcite.CalcitePlanContext;
import org.opensearch.sql.calcite.udf.mathUDF.SqrtFunction;
Expand Down Expand Up @@ -53,10 +54,28 @@ static SqlOperator translate(String op) {
case "/":
return SqlStdOperatorTable.DIVIDE;
// Built-in String Functions
case "CONCAT":
return SqlLibraryOperators.CONCAT_FUNCTION;
case "CONCAT_WS":
return SqlLibraryOperators.CONCAT_WS;
case "LIKE":
return SqlLibraryOperators.ILIKE;
case "LTRIM", "RTRIM", "TRIM":
return SqlStdOperatorTable.TRIM;
case "LENGTH":
return SqlStdOperatorTable.CHAR_LENGTH;
case "LOWER":
return SqlStdOperatorTable.LOWER;
case "LIKE":
return SqlStdOperatorTable.LIKE;
case "POSITION":
return SqlStdOperatorTable.POSITION;
case "REVERSE":
return SqlLibraryOperators.REVERSE;
case "RIGHT":
return SqlLibraryOperators.RIGHT;
case "SUBSTRING":
return SqlStdOperatorTable.SUBSTRING;
case "UPPER":
return SqlStdOperatorTable.UPPER;
// Built-in Math Functions
case "ABS":
return SqlStdOperatorTable.ABS;
Expand Down Expand Up @@ -99,6 +118,30 @@ static SqlOperator translate(String op) {
static List<RexNode> translateArgument(
String op, List<RexNode> argList, CalcitePlanContext context) {
switch (op.toUpperCase(Locale.ROOT)) {
case "TRIM":
List<RexNode> trimArgs =
new ArrayList<>(
List.of(
context.rexBuilder.makeFlag(SqlTrimFunction.Flag.BOTH),
context.rexBuilder.makeLiteral(" ")));
trimArgs.addAll(argList);
return trimArgs;
case "LTRIM":
List<RexNode> LTrimArgs =
new ArrayList<>(
List.of(
context.rexBuilder.makeFlag(SqlTrimFunction.Flag.LEADING),
context.rexBuilder.makeLiteral(" ")));
LTrimArgs.addAll(argList);
return LTrimArgs;
case "RTRIM":
List<RexNode> RTrimArgs =
new ArrayList<>(
List.of(
context.rexBuilder.makeFlag(SqlTrimFunction.Flag.TRAILING),
context.rexBuilder.makeLiteral(" ")));
RTrimArgs.addAll(argList);
return RTrimArgs;
case "ATAN":
List<RexNode> AtanArgs = new ArrayList<>(argList);
if (AtanArgs.size() == 1) {
Expand Down
239 changes: 239 additions & 0 deletions ...test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLBuiltinFunctionIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import java.io.IOException;
import org.json.JSONObject;
import org.junit.jupiter.api.Test;
import org.opensearch.client.Request;

public class CalcitePPLBuiltinFunctionIT extends CalcitePPLIntegTestCase {
@Override
Expand All @@ -31,4 +32,242 @@ public void testSqrtAndPow() {

verifyDataRows(actual, rows("Hello", 30));
}

// Test
@Test
public void testConcat() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where name=concat('He', 'llo') | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testConcatWithField() throws IOException {
Request request1 =
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/5?refresh=true");
request1.setJsonEntity(
"{\"name\":\"HelloWay\",\"age\":27,\"state\":\"Way\",\"country\":\"Canada\",\"year\":2023,\"month\":4}");
client().performRequest(request1);
JSONObject actual =
executeQuery(
String.format(
"source=%s | where name=concat('Hello', state) | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("HelloWay", 27));
}

@Test
public void testConcatWs() throws IOException {
Request request1 =
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/5?refresh=true");
request1.setJsonEntity(
"{\"name\":\"John,Way\",\"age\":27,\"state\":\"Way\",\"country\":\"Canada\",\"year\":2023,\"month\":4}");
client().performRequest(request1);
JSONObject actual =
executeQuery(
String.format(
"source=%s | where name=concat_ws(',', 'John', state) | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("John,Way", 27));
}

@Test
public void testLength() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where length(name) = 5 | fields name, age", TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testLengthShouldBeInsensitive() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where leNgTh(name) = 5 | fields name, age", TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testLower() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where lower(name) = 'hello' | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testUpper() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where upper(name) = upper('hello') | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testLike() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where like(name, '_ello%%') | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testSubstring() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where substring(name, 2, 2) = 'el' | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testPosition() {
JSONObject actual =
executeQuery(
String.format(
"source=%s | where position('He' in name) = 1 | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

@Test
public void testTrim() throws IOException {
prepareTrim();
JSONObject actual =
executeQuery(
String.format(
"source=%s | where Trim(name) = 'Jim' | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows(" Jim", 27), rows("Jim ", 57), rows(" Jim ", 70));
}

@Test
public void testRTrim() throws IOException {
prepareTrim();
JSONObject actual =
executeQuery(
String.format(
"source=%s | where RTrim(name) = 'Jim' | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Jim ", 57));
}

@Test
public void testLTrim() throws IOException {
prepareTrim();
JSONObject actual =
executeQuery(
String.format(
"source=%s | where LTrim(name) = 'Jim' | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows(" Jim", 27));
}

@Test
public void testReverse() throws IOException {
Request request1 =
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/5?refresh=true");
request1.setJsonEntity(
"{\"name\":\"DeD\",\"age\":27,\"state\":\"B.C\",\"country\":\"Canada\",\"year\":2023,\"month\":4}");
client().performRequest(request1);
JSONObject actual =
executeQuery(
String.format(
"source=%s | where Reverse(name) = name | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("DeD", 27));
}

@Test
public void testRight() throws IOException {
Request request1 =
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/5?refresh=true");
request1.setJsonEntity(
"{\"name\":\"DeD\",\"age\":27,\"state\":\"B.C\",\"country\":\"Canada\",\"year\":2023,\"month\":4}");
client().performRequest(request1);
JSONObject actual =
executeQuery(
String.format(
"source=%s | where right(name, 2) = 'lo' | fields name, age",
TEST_INDEX_STATE_COUNTRY));

verifySchema(actual, schema("name", "string"), schema("age", "integer"));

verifyDataRows(actual, rows("Hello", 30));
}

private void prepareTrim() throws IOException {
Request request1 =
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/5?refresh=true");
request1.setJsonEntity(
"{\"name\":\" "
+ " Jim\",\"age\":27,\"state\":\"B.C\",\"country\":\"Canada\",\"year\":2023,\"month\":4}");
client().performRequest(request1);
Request request2 =
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/6?refresh=true");
request2.setJsonEntity(
"{\"name\":\"Jim "
+ " \",\"age\":57,\"state\":\"B.C\",\"country\":\"Canada\",\"year\":2023,\"month\":4}");
client().performRequest(request2);
Request request3 =
new Request("PUT", "/opensearch-sql_test_index_state_country/_doc/7?refresh=true");
request3.setJsonEntity(
"{\"name\":\" Jim "
+ " \",\"age\":70,\"state\":\"B.C\",\"country\":\"Canada\",\"year\":2023,\"month\":4}");
client().performRequest(request3);
}
}

0 comments on commit b7f6cf7

Please sign in to comment.