feat(api): Add profiling support to unified query API

Enable profiling in unified query components via UnifiedQueryContext:
- profiling(boolean) builder to activate/deactivate profiling
- measure(name, callable) API for user-profiled phases (execute, format)
- getProfile() to retrieve QueryProfile snapshot after execution
- Auto-profiling in UnifiedQueryPlanner (ANALYZE), UnifiedQueryCompiler
  (OPTIMIZE), and UnifiedQueryTranspiler (TRANSPILE)
- All components use context.measure() internally for consistency
- NoopProfileContext ensures zero overhead when profiling is disabled

Signed-off-by: Chen Dai <daichen@amazon.com>

Author: dai-chen

api/README.md

@@ -62,6 +62,7 @@ Use `UnifiedQueryTranspiler` to convert Calcite logical plans into SQL strings f
 ```java
 UnifiedQueryTranspiler transpiler = UnifiedQueryTranspiler.builder()
     .dialect(SparkSqlDialect.DEFAULT)
+    .context(context)
     .build();
 
 String sql = transpiler.toSql(plan);
@@ -161,6 +162,7 @@ try (UnifiedQueryContext context = UnifiedQueryContext.builder()
   // Option A: Transpile to target SQL
   UnifiedQueryTranspiler transpiler = UnifiedQueryTranspiler.builder()
       .dialect(SparkSqlDialect.DEFAULT)
+      .context(context)
       .build();
   String sparkSql = transpiler.toSql(plan);
   // Result: SELECT * FROM `catalog`.`employees` WHERE `age` > 30
@@ -176,6 +178,67 @@ try (UnifiedQueryContext context = UnifiedQueryContext.builder()
 }
 ```
 
+## Profiling
+
+The unified query API supports the same [profiling capability](https://github.com/opensearch-project/sql/blob/main/docs/user/ppl/interfaces/endpoint.md#profile-experimental) as the PPL REST endpoint. When enabled, each unified query component automatically collects per-phase timing metrics. For code outside unified query components (e.g., `PreparedStatement.executeQuery()` or response formatting), `context.measure()` records custom phases into the same profile.
+
+```java
+try (UnifiedQueryContext context = UnifiedQueryContext.builder()
+    .language(QueryType.PPL)
+    .catalog("catalog", schema)
+    .defaultNamespace("catalog")
+    .profiling(true)
+    .build()) {
+
+  // Auto-profiled: ANALYZE
+  RelNode plan = new UnifiedQueryPlanner(context).plan(query);
+
+  // Auto-profiled: OPTIMIZE
+  PreparedStatement stmt = new UnifiedQueryCompiler(context).compile(plan);
+
+  // Auto-profiled: TRANSPILE
+  String sql = UnifiedQueryTranspiler.builder()
+      .dialect(SparkSqlDialect.DEFAULT)
+      .context(context)
+      .build()
+      .toSql(plan);
+
+  // User-profiled via measure()
+  ResultSet rs = context.measure("execute", stmt::executeQuery);
+  String json = context.measure("format", () -> formatter.format(result));
+
+  // Retrieve profile snapshot
+  QueryProfile profile = context.getProfile();
+}
+```
+
+The returned `QueryProfile` follows the same JSON structure as the REST API:
+
+```json
+{
+  "summary": {
+    "total_time_ms": 33.34
+  },
+  "phases": {
+    "analyze": { "time_ms": 8.68 },
+    "optimize": { "time_ms": 18.2 },
+    "execute": { "time_ms": 4.87 },
+    "format": { "time_ms": 0.05 },
+    "transpile": { "time_ms": 0.82 }
+  },
+  "plan": {
+    "node": "EnumerableCalc",
+    "time_ms": 4.82,
+    "rows": 2,
+    "children": [
+      { "node": "CalciteEnumerableIndexScan", "time_ms": 4.12, "rows": 2 }
+    ]
+  }
+}
+```
+
+When profiling is disabled (the default), all components execute with zero overhead.
+
 ## Development & Testing
 
 A set of unit tests is provided to validate planner behavior.

api/src/main/java/org/opensearch/sql/api/UnifiedQueryContext.java

@@ -11,8 +11,10 @@
 
 import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.Callable;
 import lombok.Value;
 import org.apache.calcite.jdbc.CalciteSchema;
 import org.apache.calcite.plan.RelTraitDef;
@@ -27,6 +29,11 @@
 import org.opensearch.sql.calcite.SysLimit;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.executor.QueryType;
+import org.opensearch.sql.monitor.profile.MetricName;
+import org.opensearch.sql.monitor.profile.ProfileContext;
+import org.opensearch.sql.monitor.profile.ProfileMetric;
+import org.opensearch.sql.monitor.profile.QueryProfile;
+import org.opensearch.sql.monitor.profile.QueryProfiling;
 
 /**
  * A reusable abstraction shared across unified query components (planner, compiler, etc.). This
@@ -42,13 +49,61 @@ public class UnifiedQueryContext implements AutoCloseable {
   /** Settings containing execution limits and feature flags used by parsers and planners. */
   Settings settings;
 
+  /**
+   * Returns the profiling result. Call after query execution to retrieve collected metrics. Returns
+   * null if profiling was not enabled.
+   */
+  public QueryProfile getProfile() {
+    ProfileContext ctx = QueryProfiling.current();
+    return ctx.isEnabled() ? ctx.finish() : null;
+  }
+
+  /**
+   * Measures the execution time of the given action and records it as a profiling metric. When
+   * profiling is disabled, the action executes with no overhead. Use this for phases outside
+   * unified query components (e.g., execution, formatting).
+   *
+   * <p>Inspired by Micrometer's {@code Timer.record()} and Dropwizard's {@code Timer.time()}.
+   *
+   * @param <T> the return type of the action
+   * @param name the metric name matching a {@link MetricName} value (case-insensitive, e.g.,
+   *     "execute", "format")
+   * @param action the action to measure
+   * @return the result of the action
+   * @throws Exception if the action throws
+   */
+  public <T> T measure(String name, Callable<T> action) throws Exception {
+    return measure(MetricName.valueOf(name.toUpperCase(Locale.ROOT)), action);
+  }
+
+  /**
+   * Measures the execution time of the given action for the specified metric. Used internally by
+   * unified query components for auto-profiling.
+   *
+   * @param <T> the return type of the action
+   * @param metricName the metric to record
+   * @param action the action to measure
+   * @return the result of the action
+   * @throws Exception if the action throws
+   */
+  public <T> T measure(MetricName metricName, Callable<T> action) throws Exception {
+    ProfileMetric metric = QueryProfiling.current().getOrCreateMetric(metricName);
+    long start = System.nanoTime();
+    try {
+      return action.call();
+    } finally {
+      metric.set(System.nanoTime() - start);
+    }
+  }
+
   /**
    * Closes the underlying resource managed by this context.
    *
    * @throws Exception if an error occurs while closing the connection
    */
   @Override
   public void close() throws Exception {
+    QueryProfiling.clear();
     if (planContext != null && planContext.connection != null) {
       planContext.connection.close();
     }
@@ -65,6 +120,7 @@ public static class Builder {
     private final Map<String, Schema> catalogs = new HashMap<>();
     private String defaultNamespace;
     private boolean cacheMetadata = false;
+    private boolean profiling = false;
 
     /**
      * Setting values with defaults from SysLimit.DEFAULT. Only includes planning-required settings
@@ -124,6 +180,18 @@ public Builder cacheMetadata(boolean cache) {
       return this;
     }
 
+    /**
+     * Enables or disables query profiling. When enabled, profiling metrics are collected during
+     * query planning and execution, retrievable via {@link UnifiedQueryContext#getProfile()}.
+     *
+     * @param enabled whether to enable profiling
+     * @return this builder instance
+     */
+    public Builder profiling(boolean enabled) {
+      this.profiling = enabled;
+      return this;
+    }
+
     /**
      * Sets a specific setting value by name.
      *
@@ -151,6 +219,7 @@ public UnifiedQueryContext build() {
       CalcitePlanContext planContext =
           CalcitePlanContext.create(
               buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
+      QueryProfiling.activate(profiling);
       return new UnifiedQueryContext(planContext, settings);
     }
 

api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java

@@ -18,6 +18,7 @@
 import org.opensearch.sql.common.antlr.Parser;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.executor.QueryType;
+import org.opensearch.sql.monitor.profile.MetricName;
 import org.opensearch.sql.ppl.antlr.PPLSyntaxParser;
 import org.opensearch.sql.ppl.parser.AstBuilder;
 import org.opensearch.sql.ppl.parser.AstStatementBuilder;
@@ -57,9 +58,8 @@ public UnifiedQueryPlanner(UnifiedQueryContext context) {
    */
   public RelNode plan(String query) {
     try {
-      return preserveCollation(analyze(parse(query)));
+      return context.measure(MetricName.ANALYZE, () -> preserveCollation(analyze(parse(query))));
     } catch (SyntaxCheckException e) {
-      // Re-throw syntax error without wrapping
       throw e;
     } catch (Exception e) {
       throw new IllegalStateException("Failed to plan query", e);

api/src/main/java/org/opensearch/sql/api/compiler/UnifiedQueryCompiler.java

@@ -16,6 +16,7 @@
 import org.apache.calcite.rel.logical.LogicalTableScan;
 import org.apache.calcite.tools.RelRunner;
 import org.opensearch.sql.api.UnifiedQueryContext;
+import org.opensearch.sql.monitor.profile.MetricName;
 
 /**
  * {@code UnifiedQueryCompiler} compiles Calcite logical plans ({@link RelNode}) into executable
@@ -46,26 +47,29 @@ public UnifiedQueryCompiler(UnifiedQueryContext context) {
    */
   public PreparedStatement compile(@NonNull RelNode plan) {
     try {
-      // Apply shuttle to convert LogicalTableScan to BindableTableScan
-      final RelHomogeneousShuttle shuttle =
-          new RelHomogeneousShuttle() {
-            @Override
-            public RelNode visit(TableScan scan) {
-              final RelOptTable table = scan.getTable();
-              if (scan instanceof LogicalTableScan
-                  && Bindables.BindableTableScan.canHandle(table)) {
-                return Bindables.BindableTableScan.create(scan.getCluster(), table);
-              }
-              return super.visit(scan);
-            }
-          };
-      RelNode transformedPlan = plan.accept(shuttle);
-
-      Connection connection = context.getPlanContext().connection;
-      final RelRunner runner = connection.unwrap(RelRunner.class);
-      return runner.prepareStatement(transformedPlan);
+      return context.measure(MetricName.OPTIMIZE, () -> doCompile(plan));
     } catch (Exception e) {
       throw new IllegalStateException("Failed to compile logical plan", e);
     }
   }
+
+  private PreparedStatement doCompile(RelNode plan) throws Exception {
+    // Apply shuttle to convert LogicalTableScan to BindableTableScan
+    final RelHomogeneousShuttle shuttle =
+        new RelHomogeneousShuttle() {
+          @Override
+          public RelNode visit(TableScan scan) {
+            final RelOptTable table = scan.getTable();
+            if (scan instanceof LogicalTableScan && Bindables.BindableTableScan.canHandle(table)) {
+              return Bindables.BindableTableScan.create(scan.getCluster(), table);
+            }
+            return super.visit(scan);
+          }
+        };
+    RelNode transformedPlan = plan.accept(shuttle);
+
+    Connection connection = context.getPlanContext().connection;
+    final RelRunner runner = connection.unwrap(RelRunner.class);
+    return runner.prepareStatement(transformedPlan);
+  }
 }

api/src/main/java/org/opensearch/sql/api/transpiler/UnifiedQueryTranspiler.java

@@ -6,10 +6,13 @@
 package org.opensearch.sql.api.transpiler;
 
 import lombok.Builder;
+import lombok.NonNull;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.rel2sql.RelToSqlConverter;
 import org.apache.calcite.sql.SqlDialect;
 import org.apache.calcite.sql.SqlNode;
+import org.opensearch.sql.api.UnifiedQueryContext;
+import org.opensearch.sql.monitor.profile.MetricName;
 
 /**
  * Transpiles Calcite logical plans ({@link RelNode}) into SQL strings for various target databases.
@@ -22,6 +25,9 @@ public class UnifiedQueryTranspiler {
   /** Target SQL dialect */
   private final SqlDialect dialect;
 
+  /** Unified query context for profiling support. */
+  @NonNull private final UnifiedQueryContext context;
+
   /**
    * Converts a Calcite logical plan to a SQL string using the configured target dialect.
    *
@@ -30,11 +36,15 @@ public class UnifiedQueryTranspiler {
    */
   public String toSql(RelNode plan) {
     try {
-      RelToSqlConverter converter = new RelToSqlConverter(dialect);
-      SqlNode sqlNode = converter.visitRoot(plan).asStatement();
-      return sqlNode.toSqlString(dialect).getSql();
+      return context.measure(MetricName.TRANSPILE, () -> doTranspile(plan));
     } catch (Exception e) {
       throw new IllegalStateException("Failed to transpile logical plan to SQL", e);
     }
   }
+
+  private String doTranspile(RelNode plan) {
+    RelToSqlConverter converter = new RelToSqlConverter(dialect);
+    SqlNode sqlNode = converter.visitRoot(plan).asStatement();
+    return sqlNode.toSqlString(dialect).getSql();
+  }
 }