Skip script encoding when run explain with 'extended' (opensearch-project#3930) (opensearch-project#4030)

Author: LantaoJin

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<Boolean> skipEncoding = ThreadLocal.withInitial(() -> false);
+
   @Getter @Setter private boolean isResolvingJoinCondition = false;
   @Getter @Setter private boolean isResolvingSubquery = false;
 

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

@@ -108,6 +108,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 {

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");

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"
+  }
+}

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();
           }
         });
   }

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);
     }

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();