diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index 7c0872b67fc..8afddd4579c 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -36,6 +36,8 @@ public class CalcitePlanContext { public final QueryType queryType; public final Integer querySizeLimit; + public static final ThreadLocal skipEncoding = ThreadLocal.withInitial(() -> false); + @Getter @Setter private boolean isResolvingJoinCondition = false; @Getter @Setter private boolean isResolvingSubquery = false; diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java index 7ab4eab104b..5fd279c5104 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java @@ -109,6 +109,17 @@ public void supportPartialPushDownScript() throws IOException { assertJsonEqualsIgnoreId(expected, result); } + @Test + public void testSkipScriptEncodingOnExtendedFormat() throws IOException { + Assume.assumeTrue("This test is only for push down enabled", isPushdownEnabled()); + String query = + "source=opensearch-sql_test_index_account | where address = '671 Bristol Street' and age -" + + " 2 = 30 | fields firstname, age, address"; + var result = explainQueryToString(query, true); + String expected = loadFromFile("expectedOutput/calcite/explain_skip_script_encoding.json"); + assertJsonEqualsIgnoreId(expected, result); + } + // Only for Calcite, as v2 gets unstable serialized string for function @Test public void testFilterScriptPushDownExplain() throws Exception { diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java index e799c3535e9..3e79d88db97 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java @@ -7,6 +7,7 @@ import static org.opensearch.sql.legacy.TestUtils.getResponseBody; import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.EXPLAIN_API_ENDPOINT; +import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.EXTENDED_EXPLAIN_API_ENDPOINT; import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.QUERY_API_ENDPOINT; import com.google.common.io.Resources; @@ -50,7 +51,15 @@ protected String executeQueryToString(String query) throws IOException { } protected String explainQueryToString(String query) throws IOException { - Response response = client().performRequest(buildRequest(query, EXPLAIN_API_ENDPOINT)); + return explainQueryToString(query, false); + } + + protected String explainQueryToString(String query, boolean extended) throws IOException { + Response response = + client() + .performRequest( + buildRequest( + query, extended ? EXTENDED_EXPLAIN_API_ENDPOINT : EXPLAIN_API_ENDPOINT)); Assert.assertEquals(200, response.getStatusLine().getStatusCode()); String responseBody = getResponseBody(response, true); return responseBody.replace("\\r\\n", "\\n"); diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_skip_script_encoding.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_skip_script_encoding.json new file mode 100644 index 00000000000..6f608d7e785 --- /dev/null +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_skip_script_encoding.json @@ -0,0 +1,7 @@ +{ + "calcite": { + "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", + "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", + "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" + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index 3334fe88bf7..4fc89d9f727 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -182,6 +182,7 @@ public void explain( try (Hook.Closeable closeable = getPhysicalPlanInHook(physical, level)) { if (format == ExplainFormat.EXTENDED) { getCodegenInHook(javaCode); + CalcitePlanContext.skipEncoding.set(true); } // triggers the hook AccessController.doPrivileged( @@ -194,6 +195,8 @@ public void explain( } } catch (Exception e) { listener.onFailure(e); + } finally { + CalcitePlanContext.skipEncoding.remove(); } }); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serde/RelJsonSerializer.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serde/RelJsonSerializer.java index 5d394942638..3c6154f23bb 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serde/RelJsonSerializer.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serde/RelJsonSerializer.java @@ -26,6 +26,7 @@ import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.util.SqlOperatorTables; import org.apache.calcite.util.JsonBuilder; +import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.function.PPLBuiltinOperators; @@ -94,7 +95,9 @@ public String serialize( ObjectOutputStream objectOutput = new ObjectOutputStream(output); objectOutput.writeObject(envelope); objectOutput.flush(); - return Base64.getEncoder().encodeToString(output.toByteArray()); + return CalcitePlanContext.skipEncoding.get() + ? rexNodeJson + : Base64.getEncoder().encodeToString(output.toByteArray()); } catch (Exception e) { throw new IllegalStateException("Failed to serialize RexNode: " + rexNode, e); } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java index 9a082b88824..53699f57fc4 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java @@ -43,6 +43,8 @@ public class RestPPLQueryAction extends BaseRestHandler { public static final String EXPLAIN_API_ENDPOINT = "/_plugins/_ppl/_explain"; public static final String LEGACY_QUERY_API_ENDPOINT = "/_opendistro/_ppl"; public static final String LEGACY_EXPLAIN_API_ENDPOINT = "/_opendistro/_ppl/_explain"; + public static final String EXTENDED_EXPLAIN_API_ENDPOINT = + "/_plugins/_ppl/_explain?format=extended"; private static final Logger LOG = LogManager.getLogger();