Skip to content

Commit 964d8b5

Browse files
LantaoJinLantaoJin
authored
Skip script encoding when run explain with 'extended' (opensearch-project#3930)
* No need to decode script when run explain command Signed-off-by: Lantao Jin <ltjin@amazon.com> * address comment Signed-off-by: Lantao Jin <ltjin@amazon.com> * Do not encoding when explain format is 'extended' Signed-off-by: Lantao Jin <ltjin@amazon.com> * Rename the thread local var Signed-off-by: Lantao Jin <ltjin@amazon.com> * Fix IT after merge main Signed-off-by: Lantao Jin <ltjin@amazon.com> --------- Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent ad3fc1f commit 964d8b5

7 files changed

Lines changed: 39 additions & 2 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ public class CalcitePlanContext {
3636
public final QueryType queryType;
3737
public final Integer querySizeLimit;
3838

39+
public static final ThreadLocal<Boolean> skipEncoding = ThreadLocal.withInitial(() -> false);
40+
3941
@Getter @Setter private boolean isResolvingJoinCondition = false;
4042
@Getter @Setter private boolean isResolvingSubquery = false;
4143

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,17 @@ public void supportPartialPushDownScript() throws IOException {
108108
assertJsonEqualsIgnoreId(expected, result);
109109
}
110110

111+
@Test
112+
public void testSkipScriptEncodingOnExtendedFormat() throws IOException {
113+
Assume.assumeTrue("This test is only for push down enabled", isPushdownEnabled());
114+
String query =
115+
"source=opensearch-sql_test_index_account | where address = '671 Bristol Street' and age -"
116+
+ " 2 = 30 | fields firstname, age, address";
117+
var result = explainQueryToString(query, true);
118+
String expected = loadFromFile("expectedOutput/calcite/explain_skip_script_encoding.json");
119+
assertJsonEqualsIgnoreId(expected, result);
120+
}
121+
111122
// Only for Calcite, as v2 gets unstable serialized string for function
112123
@Test
113124
public void testFilterScriptPushDownExplain() throws Exception {

‎integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
99
import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.EXPLAIN_API_ENDPOINT;
10+
import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.EXTENDED_EXPLAIN_API_ENDPOINT;
1011
import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.QUERY_API_ENDPOINT;
1112

1213
import com.google.common.io.Resources;
@@ -50,7 +51,15 @@ protected String executeQueryToString(String query) throws IOException {
5051
}
5152

5253
protected String explainQueryToString(String query) throws IOException {
53-
Response response = client().performRequest(buildRequest(query, EXPLAIN_API_ENDPOINT));
54+
return explainQueryToString(query, false);
55+
}
56+
57+
protected String explainQueryToString(String query, boolean extended) throws IOException {
58+
Response response =
59+
client()
60+
.performRequest(
61+
buildRequest(
62+
query, extended ? EXTENDED_EXPLAIN_API_ENDPOINT : EXPLAIN_API_ENDPOINT));
5463
Assert.assertEquals(200, response.getStatusLine().getStatusCode());
5564
String responseBody = getResponseBody(response, true);
5665
return responseBody.replace("\\r\\n", "\\n");

‎integ-test/src/test/resources/expectedOutput/calcite/explain_skip_script_encoding.json‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
{
2+
"calcite": {
3+
"logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(firstname=[$1], age=[$8], address=[$2])\n LogicalFilter(condition=[AND(=($2, '671 Bristol Street'), =(-($8, 2), 30))])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
4+
"physical": "EnumerableLimit(fetch=[10000])\n EnumerableCalc(expr#0..2=[{inputs}], expr#3=['671 Bristol Street':VARCHAR], expr#4=[=($t1, $t3)], firstname=[$t0], age=[$t2], address=[$t1], $condition=[$t4])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[firstname, address, age], SCRIPT->=(-($2, 2), 30)], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"query\":{\"bool\":{\"must\":[{\"script\":{\"script\":{\"source\":\"{\\\"langType\\\":\\\"calcite\\\",\\\"script\\\":\\\"{\\\\n \\\\\\\"op\\\\\\\": {\\\\n \\\\\\\"name\\\\\\\": \\\\\\\"=\\\\\\\",\\\\n \\\\\\\"kind\\\\\\\": \\\\\\\"EQUALS\\\\\\\",\\\\n \\\\\\\"syntax\\\\\\\": \\\\\\\"BINARY\\\\\\\"\\\\n },\\\\n \\\\\\\"operands\\\\\\\": [\\\\n {\\\\n \\\\\\\"op\\\\\\\": {\\\\n \\\\\\\"name\\\\\\\": \\\\\\\"-\\\\\\\",\\\\n \\\\\\\"kind\\\\\\\": \\\\\\\"MINUS\\\\\\\",\\\\n \\\\\\\"syntax\\\\\\\": \\\\\\\"BINARY\\\\\\\"\\\\n },\\\\n \\\\\\\"operands\\\\\\\": [\\\\n {\\\\n \\\\\\\"input\\\\\\\": 2,\\\\n \\\\\\\"name\\\\\\\": \\\\\\\"$2\\\\\\\"\\\\n },\\\\n {\\\\n \\\\\\\"literal\\\\\\\": 2,\\\\n \\\\\\\"type\\\\\\\": {\\\\n \\\\\\\"type\\\\\\\": \\\\\\\"INTEGER\\\\\\\",\\\\n \\\\\\\"nullable\\\\\\\": false\\\\n }\\\\n }\\\\n ],\\\\n \\\\\\\"type\\\\\\\": {\\\\n \\\\\\\"type\\\\\\\": \\\\\\\"BIGINT\\\\\\\",\\\\n \\\\\\\"nullable\\\\\\\": true\\\\n }\\\\n },\\\\n {\\\\n \\\\\\\"literal\\\\\\\": 30,\\\\n \\\\\\\"type\\\\\\\": {\\\\n \\\\\\\"type\\\\\\\": \\\\\\\"INTEGER\\\\\\\",\\\\n \\\\\\\"nullable\\\\\\\": false\\\\n }\\\\n }\\\\n ]\\\\n}\\\"}\",\"lang\":\"opensearch_compounded_script\",\"params\":{\"utcTimestamp\":*}},\"boost\":1.0}}],\"adjust_pure_negative\":true,\"boost\":1.0}},\"_source\":{\"includes\":[\"firstname\",\"address\",\"age\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n",
5+
"extended": "public org.apache.calcite.linq4j.Enumerable bind(final org.apache.calcite.DataContext root) {\n final org.opensearch.sql.opensearch.storage.scan.CalciteEnumerableIndexScan v1stashed = (org.opensearch.sql.opensearch.storage.scan.CalciteEnumerableIndexScan) root.get(\"v1stashed\");\n final org.apache.calcite.linq4j.Enumerable _inputEnumerable = v1stashed.scan();\n final org.apache.calcite.linq4j.AbstractEnumerable child = new org.apache.calcite.linq4j.AbstractEnumerable(){\n public org.apache.calcite.linq4j.Enumerator enumerator() {\n return new org.apache.calcite.linq4j.Enumerator(){\n public final org.apache.calcite.linq4j.Enumerator inputEnumerator = _inputEnumerable.enumerator();\n public void reset() {\n inputEnumerator.reset();\n }\n\n public boolean moveNext() {\n while (inputEnumerator.moveNext()) {\n final Object[] current = (Object[]) inputEnumerator.current();\n final String input_value = current[1] == null ? null : current[1].toString();\n final Boolean binary_call_value = input_value == null ? null : Boolean.valueOf(org.apache.calcite.runtime.SqlFunctions.eq(input_value, \"671 Bristol Street\"));\n if (binary_call_value != null && org.apache.calcite.runtime.SqlFunctions.toBoolean(binary_call_value)) {\n return true;\n }\n }\n return false;\n }\n\n public void close() {\n inputEnumerator.close();\n }\n\n public Object current() {\n final Object[] current = (Object[]) inputEnumerator.current();\n final Object input_value = current[0];\n final Object input_value0 = current[2];\n final Object input_value1 = current[1];\n return new Object[] {\n input_value,\n input_value0,\n input_value1};\n }\n\n };\n }\n\n };\n return child.take(10000);\n}\n\n\npublic Class getElementType() {\n return java.lang.Object[].class;\n}\n\n\n"
6+
}
7+
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,7 @@ public void explain(
172172
try (Hook.Closeable closeable = getPhysicalPlanInHook(physical, level)) {
173173
if (format == ExplainFormat.EXTENDED) {
174174
getCodegenInHook(javaCode);
175+
CalcitePlanContext.skipEncoding.set(true);
175176
}
176177
// triggers the hook
177178
AccessController.doPrivileged(
@@ -184,6 +185,8 @@ public void explain(
184185
}
185186
} catch (Exception e) {
186187
listener.onFailure(e);
188+
} finally {
189+
CalcitePlanContext.skipEncoding.remove();
187190
}
188191
});
189192
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serde/RelJsonSerializer.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
2727
import org.apache.calcite.sql.util.SqlOperatorTables;
2828
import org.apache.calcite.util.JsonBuilder;
29+
import org.opensearch.sql.calcite.CalcitePlanContext;
2930
import org.opensearch.sql.data.type.ExprType;
3031
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
3132

@@ -94,7 +95,9 @@ public String serialize(
9495
ObjectOutputStream objectOutput = new ObjectOutputStream(output);
9596
objectOutput.writeObject(envelope);
9697
objectOutput.flush();
97-
return Base64.getEncoder().encodeToString(output.toByteArray());
98+
return CalcitePlanContext.skipEncoding.get()
99+
? rexNodeJson
100+
: Base64.getEncoder().encodeToString(output.toByteArray());
98101
} catch (Exception e) {
99102
throw new IllegalStateException("Failed to serialize RexNode: " + rexNode, e);
100103
}

‎plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@
4141
public class RestPPLQueryAction extends BaseRestHandler {
4242
public static final String QUERY_API_ENDPOINT = "/_plugins/_ppl";
4343
public static final String EXPLAIN_API_ENDPOINT = "/_plugins/_ppl/_explain";
44+
public static final String EXTENDED_EXPLAIN_API_ENDPOINT =
45+
"/_plugins/_ppl/_explain?format=extended";
4446

4547
private static final Logger LOG = LogManager.getLogger();
4648

0 commit comments

Comments
 (0)