diff --git a/core/src/main/java/org/opensearch/sql/ast/statement/Explain.java b/core/src/main/java/org/opensearch/sql/ast/statement/Explain.java index b2fcb2b8179..01ce206ad23 100644 --- a/core/src/main/java/org/opensearch/sql/ast/statement/Explain.java +++ b/core/src/main/java/org/opensearch/sql/ast/statement/Explain.java @@ -9,6 +9,7 @@ import lombok.Getter; import org.opensearch.sql.ast.AbstractNodeVisitor; import org.opensearch.sql.executor.QueryType; +import org.opensearch.sql.protocol.response.format.Format; /** Explain Statement. */ @Getter @@ -18,15 +19,21 @@ public class Explain extends Statement { private final Statement statement; private final QueryType queryType; private final ExplainMode mode; + private final Format format; public Explain(Statement statement, QueryType queryType) { - this(statement, queryType, null); + this(statement, queryType, null, null); } public Explain(Statement statement, QueryType queryType, String mode) { + this(statement, queryType, mode, null); + } + + public Explain(Statement statement, QueryType queryType, String mode, Format format) { this.statement = statement; this.queryType = queryType; this.mode = ExplainMode.of(mode); + this.format = format; } @Override diff --git a/core/src/main/java/org/opensearch/sql/calcite/plan/rel/LogicalSystemLimit.java b/core/src/main/java/org/opensearch/sql/calcite/plan/rel/LogicalSystemLimit.java index 4f999cf0792..2ba6e8812dc 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/plan/rel/LogicalSystemLimit.java +++ b/core/src/main/java/org/opensearch/sql/calcite/plan/rel/LogicalSystemLimit.java @@ -83,8 +83,8 @@ public Sort copy( @Override public RelWriter explainTerms(RelWriter pw) { super.explainTerms(pw); - // Show type in the explain - pw.item("type", type); + // Show type in the explain - convert to string for JSON serialization compatibility + pw.item("type", type.name()); return pw; } } diff --git a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java b/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java index 2a5d392a149..9b51876c004 100644 --- a/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java +++ b/core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java @@ -21,6 +21,7 @@ import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.executor.pagination.Cursor; import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.protocol.response.format.Format; /** Execution engine that encapsulates execution details. */ public interface ExecutionEngine { @@ -75,6 +76,16 @@ default void explain( getClass().getSimpleName() + " does not support RelNode explain")); } + default void explain( + RelNode plan, + ExplainMode mode, + Format format, + CalcitePlanContext context, + ResponseListener listener) { + // Default: ignore format parameter, delegate to old signature for BWC + explain(plan, mode, context, listener); + } + /** Data class that encapsulates ExprValue. */ @Data class QueryResponse { @@ -163,5 +174,8 @@ class ExplainResponseNodeV2 { private final String logical; private final String physical; private final String extended; + // For json_tree format: parsed JSON objects instead of strings + private Object logicalTree; + private Object physicalTree; } } diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index fe9d3e55dc1..ddb4338bc8f 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -54,6 +54,7 @@ import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.protocol.response.format.Format; /** The low level interface of core engine. */ @RequiredArgsConstructor @@ -124,8 +125,19 @@ public void explain( HighlightConfig highlightConfig, ResponseListener listener, ExplainMode mode) { + explain(plan, queryType, highlightConfig, listener, mode, null); + } + + /** Explain with optional highlight config and format. */ + public void explain( + UnresolvedPlan plan, + QueryType queryType, + HighlightConfig highlightConfig, + ResponseListener listener, + ExplainMode mode, + Format format) { if (shouldUseCalcite(queryType)) { - explainWithCalcite(plan, queryType, highlightConfig, listener, mode); + explainWithCalcite(plan, queryType, highlightConfig, listener, mode, format); } else { explainWithLegacy(plan, queryType, listener, mode, Optional.empty()); } @@ -192,6 +204,16 @@ public void explainWithCalcite( HighlightConfig highlightConfig, ResponseListener listener, ExplainMode mode) { + explainWithCalcite(plan, queryType, highlightConfig, listener, mode, null); + } + + public void explainWithCalcite( + UnresolvedPlan plan, + QueryType queryType, + HighlightConfig highlightConfig, + ResponseListener listener, + ExplainMode mode, + Format format) { CalcitePlanContext.run( () -> { try { @@ -206,7 +228,11 @@ public void explainWithCalcite( () -> { RelNode relNode = analyze(plan, context); RelNode calcitePlan = convertToCalcitePlan(relNode, context); - executionEngine.explain(calcitePlan, mode, context, listener); + if (format != null) { + executionEngine.explain(calcitePlan, mode, format, context, listener); + } else { + executionEngine.explain(calcitePlan, mode, context, listener); + } }, settings); }, diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java index e470d12507e..fbdabe2fa44 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java @@ -12,6 +12,7 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryType; +import org.opensearch.sql.protocol.response.format.Format; /** AbstractPlan represent the execution entity of the Statement. */ @RequiredArgsConstructor @@ -32,4 +33,16 @@ public abstract class AbstractPlan { */ public abstract void explain( ResponseListener listener, ExplainMode mode); + + /** + * Explain query execution with format. + * + * @param listener query explain response listener. + * @param mode explain mode + * @param format output format + */ + public void explain( + ResponseListener listener, ExplainMode mode, Format format) { + explain(listener, mode); + } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/ExplainPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/ExplainPlan.java index 27f7a47e504..0a196c6f484 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/ExplainPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/ExplainPlan.java @@ -10,12 +10,14 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryType; +import org.opensearch.sql.protocol.response.format.Format; /** Explain plan. */ public class ExplainPlan extends AbstractPlan { private final AbstractPlan plan; private final ExplainMode mode; + private final Format format; private final ResponseListener explainListener; @@ -26,15 +28,27 @@ public ExplainPlan( AbstractPlan plan, ExplainMode mode, ResponseListener explainListener) { + this(queryId, queryType, plan, mode, null, explainListener); + } + + /** Constructor with format. */ + public ExplainPlan( + QueryId queryId, + QueryType queryType, + AbstractPlan plan, + ExplainMode mode, + Format format, + ResponseListener explainListener) { super(queryId, queryType); this.plan = plan; this.mode = mode; + this.format = format; this.explainListener = explainListener; } @Override public void execute() { - plan.explain(explainListener, mode); + plan.explain(explainListener, mode, format); } @Override @@ -42,4 +56,10 @@ public void explain( ResponseListener listener, ExplainMode mode) { throw new UnsupportedOperationException("explain query can not been explained."); } + + @Override + public void explain( + ResponseListener listener, ExplainMode mode, Format format) { + throw new UnsupportedOperationException("explain query can not been explained."); + } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java index 3f6407e8873..762997439f4 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java @@ -16,6 +16,7 @@ import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.QueryType; +import org.opensearch.sql.protocol.response.format.Format; /** Query plan which includes a select query. */ public class QueryPlan extends AbstractPlan { @@ -86,12 +87,18 @@ public void execute() { @Override public void explain( ResponseListener listener, ExplainMode mode) { + explain(listener, mode, null); + } + + @Override + public void explain( + ResponseListener listener, ExplainMode mode, Format format) { if (pageSize.isPresent()) { listener.onFailure( new NotImplementedException( "`explain` feature for paginated requests is not implemented yet.")); } else { - queryService.explain(plan, getQueryType(), highlightConfig, listener, mode); + queryService.explain(plan, getQueryType(), highlightConfig, listener, mode, format); } } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 48e2b3ce5e0..93c73a2315b 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -144,6 +144,7 @@ public AbstractPlan visitExplain( node.getQueryType(), create(node.getStatement(), NO_CONSUMER_RESPONSE_LISTENER, context.getRight()), node.getMode(), + node.getFormat(), context.getRight()); } } diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/Format.java b/core/src/main/java/org/opensearch/sql/protocol/response/format/Format.java similarity index 93% rename from protocol/src/main/java/org/opensearch/sql/protocol/response/format/Format.java rename to core/src/main/java/org/opensearch/sql/protocol/response/format/Format.java index 6db28e4d4ab..b8daf3bbd85 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/Format.java +++ b/core/src/main/java/org/opensearch/sql/protocol/response/format/Format.java @@ -25,6 +25,8 @@ public enum Format { JSON("json"), /** Returns explain output in yaml format */ YAML("yaml"), + /** Returns explain output as structured JSON tree using RelJsonWriter */ + JSON_TREE("json_tree"), /*---- backward compatible format of explain response -----*/ SIMPLE("simple"), @@ -52,6 +54,7 @@ public enum Format { builder = new ImmutableMap.Builder<>(); builder.put(JSON.formatName, JSON); builder.put(YAML.formatName, YAML); + builder.put(JSON_TREE.formatName, JSON_TREE); builder.put(SIMPLE.formatName, SIMPLE); builder.put(STANDARD.formatName, STANDARD); builder.put(EXTENDED.formatName, EXTENDED); diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/ExplainPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/ExplainPlanTest.java index 4cb5c755d14..977c1cb4729 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/ExplainPlanTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/ExplainPlanTest.java @@ -36,12 +36,12 @@ public class ExplainPlanTest { @Test public void execute() { - doNothing().when(queryPlan).explain(any(), any()); + doNothing().when(queryPlan).explain(any(), any(), any()); ExplainPlan explainPlan = new ExplainPlan(queryId, queryType, queryPlan, mode, explainListener); explainPlan.execute(); - verify(queryPlan, times(1)).explain(explainListener, mode); + verify(queryPlan, times(1)).explain(explainListener, mode, null); } @Test diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java index 128df14ff8e..6e05c8c1258 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java @@ -58,9 +58,9 @@ public void execute_no_page_size() { @Test public void explain_no_page_size() { QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener); - query.explain(explainListener, mode); + query.explain(explainListener, mode, null); - verify(queryService, times(1)).explain(plan, queryType, null, explainListener, mode); + verify(queryService, times(1)).explain(plan, queryType, null, explainListener, mode, null); } @Test 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 95c47b9b0b7..d475f11427d 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 @@ -28,6 +28,8 @@ import java.io.IOException; import java.util.Locale; import org.apache.commons.text.StringEscapeUtils; +import org.json.JSONArray; +import org.json.JSONObject; import org.junit.Assert; import org.junit.Ignore; import org.junit.Test; @@ -3010,4 +3012,29 @@ public void testExplainUnion() throws IOException { String expected = loadExpectedPlan("explain_union.yaml"); assertYamlEqualsIgnoreId(expected, actual); } + + @Test + public void testExplainJsonTreeFormat() throws IOException { + String query = "source=opensearch-sql_test_index_account | where age > 30 | fields age"; + String result = explainQuery(query, Format.JSON_TREE, ExplainMode.STANDARD); + + // Parse JSON response + JSONObject json = new JSONObject(result); + JSONObject calcite = json.getJSONObject("calcite"); + + // Verify logical plan is a structured JSON object (not a plain string) + JSONObject logical = calcite.getJSONObject("logical"); + Assert.assertTrue("Logical plan should contain 'rels' array", logical.has("rels")); + JSONArray rels = logical.getJSONArray("rels"); + Assert.assertTrue("Rels array should not be empty", rels.length() > 0); + + // Verify first rel is a proper RelNode structure + JSONObject firstRel = rels.getJSONObject(0); + Assert.assertTrue("RelNode should have 'relOp' field", firstRel.has("relOp")); + Assert.assertTrue("RelNode should have 'id' field", firstRel.has("id")); + + // Verify physical plan also has structured format + JSONObject physical = calcite.getJSONObject("physical"); + Assert.assertTrue("Physical plan should contain 'rels' array", physical.has("rels")); + } } diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLExplainIT.java index 674a7d96f8d..78ff6fc0401 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLExplainIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLExplainIT.java @@ -7,10 +7,14 @@ import static org.opensearch.sql.util.MatcherUtils.assertJsonEquals; +import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; import org.junit.jupiter.api.Test; import org.opensearch.client.Request; +import org.opensearch.sql.ast.statement.ExplainMode; import org.opensearch.sql.ppl.PPLIntegTestCase; +import org.opensearch.sql.ppl.PPLIntegTestCase.GlobalPushdownConfig; +import org.opensearch.sql.protocol.response.format.Format; public class CalcitePPLExplainIT extends PPLIntegTestCase { @@ -74,6 +78,51 @@ public void testExplainCommandSimple() throws IOException { assertJsonEquals(expected, result); } + @Test + public void testJsonTreeFormat() throws IOException { + var resultStr = + explainQuery( + "source=test | where age > 20 | fields name", Format.JSON_TREE, ExplainMode.STANDARD); + + // Parse JSON + var mapper = new ObjectMapper(); + var result = mapper.readTree(resultStr); + + // Verify tree structure exists + assertTrue(result.has("calcite")); + assertTrue(result.get("calcite").has("logical")); + assertTrue(result.get("calcite").has("physical")); + + // Verify logical and physical are parsed JSON objects, not strings + assertTrue(result.get("calcite").get("logical").isObject()); + assertTrue(result.get("calcite").get("physical").isObject()); + + // Verify sourceBuilder exists in physical plan rels + var physical = result.get("calcite").get("physical"); + assertTrue(physical.has("rels")); + var rels = physical.get("rels"); + assertTrue(rels.isArray()); + + // Find a rel with sourceBuilder (only present when pushdown is enabled) + boolean foundSourceBuilder = false; + for (int i = 0; i < rels.size(); i++) { + var rel = rels.get(i); + if (rel.has("sourceBuilder")) { + foundSourceBuilder = true; + // Verify sourceBuilder is a parsed JSON object, not a string + assertTrue(rel.get("sourceBuilder").isObject()); + // Verify it has expected OpenSearch DSL fields + assertTrue(rel.get("sourceBuilder").has("from")); + assertTrue(rel.get("sourceBuilder").has("size")); + break; + } + } + // Only assert sourceBuilder exists when pushdown is enabled + if (GlobalPushdownConfig.enabled) { + assertTrue("sourceBuilder not found in physical plan rels", foundSourceBuilder); + } + } + /** * Executes the PPL query and returns the result as a string with windows-style line breaks * replaced with Unix-style ones. 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 3db2142cffe..f22cea74734 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 @@ -108,7 +108,7 @@ protected String explainQueryToString(String query, ExplainMode mode) throws IOE return explainQuery(query, Format.JSON, mode).replace("\\r\\n", "\\n"); } - private String explainQuery(String query, Format format, ExplainMode mode) throws IOException { + protected String explainQuery(String query, Format format, ExplainMode mode) throws IOException { Response response = client() .performRequest(buildRequest(query, String.format(EXPLAIN_API_ENDPOINT, format, mode))); 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 32b7891d344..3a1fa9fe78d 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 @@ -5,6 +5,7 @@ package org.opensearch.sql.opensearch.executor; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Suppliers; import java.sql.PreparedStatement; import java.sql.ResultSet; @@ -24,6 +25,7 @@ import org.apache.calcite.plan.RelOptUtil; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelRoot; +import org.apache.calcite.rel.externalize.RelJsonWriter; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeField; import org.apache.calcite.runtime.Hook; @@ -65,12 +67,14 @@ import org.opensearch.sql.opensearch.functions.DistinctCountApproxAggFunction; import org.opensearch.sql.opensearch.functions.GeoIpFunction; import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.protocol.response.format.Format; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.transport.client.node.NodeClient; /** OpenSearch execution engine implementation. */ public class OpenSearchExecutionEngine implements ExecutionEngine { private static final Logger logger = LogManager.getLogger(OpenSearchExecutionEngine.class); + private static final ObjectMapper objectMapper = new ObjectMapper(); private final OpenSearchClient client; @@ -167,38 +171,146 @@ private Hook.Closeable getCodegenInHook(AtomicReference codegen) { }); } + /** + * Parse sourceBuilder JSON strings within the physical plan tree to objects. This finds any + * sourceBuilder fields (which are serialized as JSON strings by RelJsonWriter) and parses them to + * JSON objects for easier client consumption. + */ + @SuppressWarnings("unchecked") + private void parseSourceBuilderInPhysicalTree(Object physicalTree) { + try { + if (!(physicalTree instanceof Map)) { + return; + } + Map tree = (Map) physicalTree; + Object relsObj = tree.get("rels"); + if (!(relsObj instanceof List)) { + return; + } + + List rels = (List) relsObj; + for (Object relObj : rels) { + if (!(relObj instanceof Map)) { + continue; + } + Map rel = (Map) relObj; + + // Parse sourceBuilder if it exists as a JSON string + Object sourceBuilderObj = rel.get("sourceBuilder"); + if (sourceBuilderObj instanceof String) { + try { + String sourceBuilderJson = (String) sourceBuilderObj; + Object parsed = objectMapper.readValue(sourceBuilderJson, Object.class); + rel.put("sourceBuilder", parsed); + } catch (Exception e) { + logger.debug("Failed to parse sourceBuilder JSON: {}", e.getMessage()); + } + } + } + } catch (Exception e) { + logger.warn("Failed to parse sourceBuilder in physical tree: " + e.getMessage()); + } + } + @Override public void explain( RelNode rel, ExplainMode mode, CalcitePlanContext context, ResponseListener listener) { + explain(rel, mode, null, context, listener); + } + + @Override + public void explain( + RelNode rel, + ExplainMode mode, + Format format, + CalcitePlanContext context, + ResponseListener listener) { client.schedule( () -> { try { - if (mode == ExplainMode.SIMPLE) { - String logical = RelOptUtil.toString(rel, SqlExplainLevel.NO_ATTRIBUTES); - listener.onResponse( - new ExplainResponse(new ExplainResponseNodeV2(logical, null, null))); + if (format == Format.JSON_TREE) { + // Use RelJsonWriter for structured JSON tree output + try { + RelJsonWriter logicalWriter = new RelJsonWriter(); + rel.explain(logicalWriter); + String logicalJson = logicalWriter.asString(); + + AtomicReference physicalJson = new AtomicReference<>(); + AtomicReference physicalError = new AtomicReference<>(); + SqlExplainLevel level = + mode == ExplainMode.COST + ? SqlExplainLevel.ALL_ATTRIBUTES + : SqlExplainLevel.EXPPLAN_ATTRIBUTES; + + try (Hook.Closeable closeable = + Hook.PLAN_BEFORE_IMPLEMENTATION.addThread( + obj -> { + try { + RelRoot relRoot = (RelRoot) obj; + RelJsonWriter physicalWriter = new RelJsonWriter(); + relRoot.rel.explain(physicalWriter); + physicalJson.set(physicalWriter.asString()); + } catch (Exception e) { + physicalError.set(e); + } + })) { + // triggers the hook + OpenSearchRelRunners.run(context, rel); + } + + if (physicalError.get() != null) { + throw physicalError.get(); + } + + // Parse JSON strings to objects for structured output + Object logicalTree = objectMapper.readValue(logicalJson, Object.class); + Object physicalTree = objectMapper.readValue(physicalJson.get(), Object.class); + + // Parse sourceBuilder JSON if present in physical plan + parseSourceBuilderInPhysicalTree(physicalTree); + + ExplainResponseNodeV2 response = + new ExplainResponseNodeV2(logicalJson, physicalJson.get(), null); + response.setLogicalTree(logicalTree); + response.setPhysicalTree(physicalTree); + + listener.onResponse(new ExplainResponse(response)); + } catch (Exception e) { + // RelJsonWriter can't handle some custom types (e.g., SystemLimitType enum) + listener.onFailure( + new UnsupportedOperationException( + "Cannot serialize plan to json_tree format: " + e.getMessage(), e)); + return; + } } else { - SqlExplainLevel level = - mode == ExplainMode.COST - ? SqlExplainLevel.ALL_ATTRIBUTES - : SqlExplainLevel.EXPPLAN_ATTRIBUTES; - String logical = RelOptUtil.toString(rel, level); - AtomicReference physical = new AtomicReference<>(); - AtomicReference javaCode = new AtomicReference<>(); - try (Hook.Closeable closeable = getPhysicalPlanInHook(physical, level)) { - if (mode == ExplainMode.EXTENDED) { - getCodegenInHook(javaCode); - CalcitePlanContext.skipEncoding.set(true); + // Original string format for json/yaml + if (mode == ExplainMode.SIMPLE) { + String logical = RelOptUtil.toString(rel, SqlExplainLevel.NO_ATTRIBUTES); + listener.onResponse( + new ExplainResponse(new ExplainResponseNodeV2(logical, null, null))); + } else { + SqlExplainLevel level = + mode == ExplainMode.COST + ? SqlExplainLevel.ALL_ATTRIBUTES + : SqlExplainLevel.EXPPLAN_ATTRIBUTES; + String logical = RelOptUtil.toString(rel, level); + AtomicReference physical = new AtomicReference<>(); + AtomicReference javaCode = new AtomicReference<>(); + try (Hook.Closeable closeable = getPhysicalPlanInHook(physical, level)) { + if (mode == ExplainMode.EXTENDED) { + getCodegenInHook(javaCode); + CalcitePlanContext.skipEncoding.set(true); + } + // triggers the hook + OpenSearchRelRunners.run(context, rel); } - // triggers the hook - OpenSearchRelRunners.run(context, rel); + listener.onResponse( + new ExplainResponse( + new ExplainResponseNodeV2(logical, physical.get(), javaCode.get()))); } - listener.onResponse( - new ExplainResponse( - new ExplainResponseNodeV2(logical, physical.get(), javaCode.get()))); } } catch (Exception e) { listener.onFailure(e); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java index 609a5aaa92f..699d15f9ed7 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java @@ -33,6 +33,7 @@ import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.externalize.RelJsonWriter; import org.apache.calcite.rel.externalize.RelWriterImpl; import org.apache.calcite.rel.hint.RelHint; import org.apache.calcite.rel.logical.LogicalAggregate; @@ -105,13 +106,29 @@ public RelDataType deriveRowType() { @Override public RelWriter explainTerms(RelWriter pw) { + // Build explain string with context and request builder info String explainString = String.valueOf(pushDownContext); - if (pw instanceof RelWriterImpl) { - // Only add request builder to the explain plan - explainString += ", " + pushDownContext.createRequestBuilder(); + if (pw instanceof RelJsonWriter) { + // For JSON output, add structured items + super.explainTerms(pw); + if (!pushDownContext.isEmpty()) { + pw.item("PushDownContext", explainString); + try { + OpenSearchRequestBuilder requestBuilder = pushDownContext.createRequestBuilder(); + pw.item("sourceBuilder", requestBuilder.getSourceBuilder().toString()); + } catch (Exception e) { + // Ignore if request builder cannot be created + } + } + return pw; + } else { + // For text output, use original chained format + if (pw instanceof RelWriterImpl && !pushDownContext.isEmpty()) { + explainString += ", " + pushDownContext.createRequestBuilder(); + } + return super.explainTerms(pw) + .itemIf("PushDownContext", explainString, !pushDownContext.isEmpty()); } - return super.explainTerms(pw) - .itemIf("PushDownContext", explainString, !pushDownContext.isEmpty()); } protected Integer getQuerySizeLimit() { diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java index 973f00c54cb..678ed58f37f 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java @@ -10,7 +10,9 @@ import static org.opensearch.sql.lang.PPLLangSpec.PPL_SPEC; import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY; +import java.util.LinkedHashMap; import java.util.Locale; +import java.util.Map; import java.util.Optional; import java.util.function.Supplier; import org.apache.calcite.rel.RelNode; @@ -243,6 +245,18 @@ protected Object buildYamlObject(ExecutionEngine.ExplainResponse response) { new JsonResponseFormatter<>(PRETTY) { @Override protected Object buildJsonObject(ExecutionEngine.ExplainResponse response) { + // For json_tree format, use parsed tree objects instead of strings + if (response.getCalcite() != null + && response.getCalcite().getLogicalTree() != null) { + Map result = new LinkedHashMap<>(); + Map calcite = new LinkedHashMap<>(); + calcite.put("logical", response.getCalcite().getLogicalTree()); + if (response.getCalcite().getPhysicalTree() != null) { + calcite.put("physical", response.getCalcite().getPhysicalTree()); + } + result.put("calcite", calcite); + return result; + } return response; } }; diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java index ea353066cb0..6ad9032432c 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -99,7 +99,12 @@ private AbstractPlan plan( .isExplain(request.isExplainRequest()) .fetchSize(request.getFetchSize()) .highlightConfig(request.getHighlightConfig()) - .format(request.getFormat()) + .format( + request.getFormat() != null && !request.getFormat().isEmpty() + ? org.opensearch.sql.protocol.response.format.Format.ofExplain( + request.getFormat()) + .orElse(null) + : null) .explainMode(request.getExplainMode()) .build())); diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java index d2c1f610238..62503923eee 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstStatementBuilder.java @@ -21,6 +21,7 @@ import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser; import org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParserBaseVisitor; +import org.opensearch.sql.protocol.response.format.Format; /** Build {@link Statement} from PPL Query. */ @RequiredArgsConstructor @@ -45,12 +46,15 @@ public Statement visitPplStatement(OpenSearchPPLParser.PplStatementContext ctx) } if (ctx.explainStatement() != null) { if (ctx.explainStatement().explainMode() == null) { - return new Explain(query, PPL); + return new Explain(query, PPL, null, context.format); } else { - return new Explain(query, PPL, ctx.explainStatement().explainMode().getText()); + return new Explain( + query, PPL, ctx.explainStatement().explainMode().getText(), context.format); } } else { - return context.isExplain ? new Explain(query, PPL, context.explainMode) : query; + return context.isExplain + ? new Explain(query, PPL, context.explainMode, context.format) + : query; } } @@ -74,7 +78,7 @@ public static class StatementBuilderContext { /** Highlight config from the API request. */ private final HighlightConfig highlightConfig; - private final String format; + private final Format format; private final String explainMode; } diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index 0825f6d1def..37398bbf7e8 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -22,7 +22,6 @@ import org.opensearch.sql.executor.DefaultQueryManager; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; -import org.opensearch.sql.executor.ExecutionEngine.ExplainResponseNode; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; @@ -134,15 +133,6 @@ public void testExecuteCsvFormatShouldPass() { @Test public void testExplainShouldPass() { - doAnswer( - invocation -> { - ResponseListener listener = invocation.getArgument(3); - listener.onResponse(new ExplainResponse(new ExplainResponseNode("test"))); - return null; - }) - .when(queryService) - .explain(any(), any(), any(), any(), any()); - pplService.explain( new PPLQueryRequest("search source=t a=1", null, EXPLAIN), new ResponseListener() {