Add explain check

Signed-off-by: Tomoyuki Morita <[email protected]>

Author: ykmr1224

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

@@ -30,33 +30,67 @@ public void init() throws IOException {
 
   @Test
   public void testBasicProjection() throws IOException {
-    JSONObject result =
-        executeQuery(
-            source(TEST_INDEX_DYNAMIC, "fields firstname, lastname, department, salary | head 1"));
-
+    String query =
+        source(TEST_INDEX_DYNAMIC, "fields firstname, lastname, department, salary | head 1");
+    JSONObject result = executeQuery(query);
     verifySchema(
         result,
         schema("firstname", "string"),
         schema("lastname", "string"),
         schema("department", "string"),
         schema("salary", "int"));
     verifyDataRows(result, rows("John", "Doe", "Engineering", 75000));
+
+    assertExplainYaml(
+        query,
+        "calcite:\n"
+            + "  logical: |\n"
+            + "    LogicalSystemLimit(fetch=[200], type=[QUERY_SIZE_LIMIT])\n"
+            + "      LogicalSort(fetch=[1])\n"
+            + "        LogicalProject(firstname=[$0], lastname=[$2], department=[ITEM($9,"
+            + " 'department')], salary=[ITEM($9, 'salary')])\n"
+            + "          CalciteLogicalIndexScan(table=[[OpenSearch, test_dynamic_fields]])\n"
+            + "  physical: |\n"
+            + "    EnumerableLimit(fetch=[200])\n"
+            + "      EnumerableCalc(expr#0..9=[{inputs}], expr#10=['department'],"
+            + " expr#11=[ITEM($t9, $t10)], expr#12=['salary'], expr#13=[ITEM($t9, $t12)],"
+            + " firstname=[$t0], lastname=[$t2], department=[$t11], salary=[$t13])\n"
+            + "        EnumerableLimit(fetch=[1])\n"
+            + "          CalciteEnumerableIndexScan(table=[[OpenSearch, test_dynamic_fields]])\n");
   }
 
   @Test
   public void testEval() throws IOException {
-    JSONObject result =
-        executeQuery(
-            source(
-                TEST_INDEX_DYNAMIC,
-                "eval salary = cast(salary as int) * 2 | fields firstname,"
-                    + " lastname, salary | head 1"));
+    String query =
+        source(
+            TEST_INDEX_DYNAMIC,
+            "eval salary = cast(salary as int) * 2 | fields firstname,"
+                + " lastname, salary | head 1");
+    JSONObject result = executeQuery(query);
 
     verifySchema(
         result,
         schema("firstname", "string"),
         schema("lastname", "string"),
         schema("salary", "int"));
+
+    assertExplainYaml(
+        query,
+        "calcite:\n"
+            + "  logical: |\n"
+            + "    LogicalSystemLimit(fetch=[200], type=[QUERY_SIZE_LIMIT])\n"
+            + "      LogicalSort(fetch=[1])\n"
+            + "        LogicalProject(firstname=[$0], lastname=[$2], salary=[*(SAFE_CAST(ITEM($9,"
+            + " 'salary')), 2)])\n"
+            + "          CalciteLogicalIndexScan(table=[[OpenSearch, test_dynamic_fields]])\n"
+            + "  physical: |\n"
+            + "    EnumerableLimit(fetch=[200])\n"
+            + "      EnumerableCalc(expr#0..9=[{inputs}], expr#10=['salary'], expr#11=[ITEM($t9,"
+            + " $t10)], expr#12=[SAFE_CAST($t11)], expr#13=[2], expr#14=[*($t12, $t13)],"
+            + " firstname=[$t0], lastname=[$t2], salary=[$t14])\n"
+            + "        EnumerableLimit(fetch=[1])\n"
+            + "          CalciteEnumerableIndexScan(table=[[OpenSearch, test_dynamic_fields]])\n"
+            + "");
   }
 
   private void createTestIndexWithUnmappedFields() throws IOException {
@@ -67,8 +101,8 @@ private void createTestIndexWithUnmappedFields() throws IOException {
     String mapping =
         "{"
             + "\"mappings\": {"
-            + "  \"dynamic\": false," // Disable dynamic mapping - extra fields won't be indexed but
-            // will be stored
+            // Disable dynamic mapping - extra fields won't be indexed but will be stored
+            + "  \"dynamic\": false,"
             + "  \"properties\": {"
             + "    \"firstname\": {\"type\": \"text\"},"
             + "    \"lastname\": {\"type\": \"text\"},"

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

@@ -65,7 +65,9 @@
 import org.opensearch.sql.ppl.domain.PPLQueryRequest;
 import org.opensearch.sql.protocol.response.QueryResult;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
+import org.opensearch.sql.protocol.response.format.ResponseFormatter;
 import org.opensearch.sql.protocol.response.format.SimpleJsonResponseFormatter;
+import org.opensearch.sql.protocol.response.format.YamlResponseFormatter;
 import org.opensearch.sql.sql.SQLService;
 import org.opensearch.sql.sql.antlr.SQLSyntaxParser;
 import org.opensearch.sql.storage.DataSourceFactory;
@@ -232,29 +234,50 @@ public void onFailure(Exception e) {
   protected String explainQuery(String query) {
     AtomicReference<String> actual = new AtomicReference<>();
     pplService.explain(
-        new PPLQueryRequest(query, null, null),
-        new ResponseListener<ExecutionEngine.ExplainResponse>() {
-
-          @Override
-          public void onResponse(ExecutionEngine.ExplainResponse response) {
-            String responseContent =
-                new JsonResponseFormatter<ExecutionEngine.ExplainResponse>(PRETTY) {
-                  @Override
-                  protected Object buildJsonObject(ExecutionEngine.ExplainResponse response) {
-                    return response;
-                  }
-                }.format(response);
-            actual.set(responseContent.replace("\\r\\n", "\\n"));
-          }
+        new PPLQueryRequest("explain " + query, null, null), formatListener(jsonFormatter, actual));
+    return actual.get();
+  }
 
-          @Override
-          public void onFailure(Exception e) {
-            throw new IllegalStateException("Exception happened during execution", e);
-          }
-        });
+  @Override
+  protected String explainQueryYaml(String query) {
+    AtomicReference<String> actual = new AtomicReference<>();
+    pplService.explain(
+        new PPLQueryRequest("explain " + query, null, null), formatListener(yamlFormatter, actual));
     return actual.get();
   }
 
+  private <T> ResponseListener<T> formatListener(
+      ResponseFormatter<T> formatter, AtomicReference<String> actual) {
+    return new ResponseListener<T>() {
+      @Override
+      public void onResponse(T response) {
+        String responseContent = formatter.format(response);
+        actual.set(responseContent.replace("\\r\\n", "\\n"));
+      }
+
+      @Override
+      public void onFailure(Exception e) {
+        throw new IllegalStateException("Exception happened during execution", e);
+      }
+    };
+  }
+
+  private JsonResponseFormatter<ExecutionEngine.ExplainResponse> jsonFormatter =
+      new JsonResponseFormatter<>(PRETTY) {
+        @Override
+        protected Object buildJsonObject(ExecutionEngine.ExplainResponse response) {
+          return response;
+        }
+      };
+
+  private YamlResponseFormatter<ExecutionEngine.ExplainResponse> yamlFormatter =
+      new YamlResponseFormatter<>() {
+        @Override
+        protected Object buildYamlObject(ExecutionEngine.ExplainResponse response) {
+          return response;
+        }
+      };
+
   public static DataSourceMetadataStorage getDataSourceMetadataStorage() {
     return new DataSourceMetadataStorage() {
       @Override

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

@@ -34,4 +34,19 @@ public DataSource createDataSource(DataSourceMetadata metadata) {
   protected Settings getSettings() {
     return defaultSettings(getDefaultSettingsBuilder().put(Key.PPL_QUERY_PERMISSIVE, true).build());
   }
+
+  protected void assertExplainYaml(String query, String expectedYaml) {
+    String actualYaml = explainQueryYaml(query);
+    assertTrue(getMessage(expectedYaml, actualYaml), expectedYaml.equals(actualYaml));
+  }
+
+  private String getMessage(String expectedYaml, String actualYaml) {
+    return "Explain did not match:\n"
+        + String.format(
+            "# Expected: %s# Actual: %s", blockQuote(expectedYaml), blockQuote(actualYaml));
+  }
+
+  private String blockQuote(String str) {
+    return "```\n" + str + "```\n";
+  }
 }