add more tests

Signed-off-by: Isha Gupta <igupta24@apple.com>

Author: Isha Gupta

.github/workflows/analytics-engine-compat.yml

@@ -0,0 +1,44 @@
+name: Analytics Engine Compatibility
+
+on:
+  pull_request:
+  push:
+    branches-ignore:
+      - 'backport/**'
+      - 'dependabot/**'
+    paths:
+      - '**/*.java'
+      - '**gradle*'
+      - 'integ-test/**'
+      - '.github/workflows/analytics-engine-compat.yml'
+  merge_group:
+
+jobs:
+  Get-CI-Image-Tag:
+    uses: opensearch-project/opensearch-build/.github/workflows/get-ci-image-tag.yml@main
+    with:
+      product: opensearch
+
+  analytics-engine-compat:
+    needs: Get-CI-Image-Tag
+    runs-on: ubuntu-latest
+    container:
+      image: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-version-linux }}
+      options: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-options }}
+
+    steps:
+      - name: Run start commands
+        run: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-command }}
+
+      - uses: actions/checkout@v4
+
+      - name: Set up JDK 25
+        uses: actions/setup-java@v4
+        with:
+          distribution: 'temurin'
+          java-version: 25
+
+      - name: Run analytics-engine compatibility smoke test
+        run: |
+          chown -R 1000:1000 `pwd`
+          su `id -un 1000` -c "./gradlew :integ-test:analyticsEngineCompatIT"

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

@@ -5,7 +5,9 @@
 
 package org.opensearch.sql.api;
 
+import static org.opensearch.sql.common.setting.Settings.Key.CALCITE_ENGINE_ENABLED;
 import static org.opensearch.sql.common.setting.Settings.Key.PPL_JOIN_SUBSEARCH_MAXOUT;
+import static org.opensearch.sql.common.setting.Settings.Key.PPL_REX_MAX_MATCH_LIMIT;
 import static org.opensearch.sql.common.setting.Settings.Key.PPL_SUBSEARCH_MAXOUT;
 import static org.opensearch.sql.common.setting.Settings.Key.QUERY_SIZE_LIMIT;
 
@@ -119,13 +121,31 @@ public static class Builder {
     /**
      * Setting values with defaults from SysLimit.DEFAULT. Only includes planning-required settings
      * to avoid coupling with OpenSearchSettings.
+     *
+     * <p>{@link Settings.Key#CALCITE_ENGINE_ENABLED} defaults to {@code true} here because the
+     * unified query path is by definition Calcite-based — every query reaching this context flows
+     * through Calcite's planner, never the v2 engine. The PPL {@link
+     * org.opensearch.sql.api.parser.PPLQueryParser} reuses the v2 {@code AstBuilder}, which gates
+     * Calcite-only commands (e.g. {@code visitTableCommand}) on this setting; without the default,
+     * those commands fail at parse time even when the cluster setting is true.
+     *
+     * <p>{@link Settings.Key#PPL_REX_MAX_MATCH_LIMIT} defaults to {@code 10} here because {@code
+     * AstBuilder.visitRexCommand} reads it unconditionally and unboxes to {@code int} — a {@code
+     * null} return from {@code getSettingValue} NPEs the planner before any operator-level
+     * capability check runs. The value mirrors the cluster-side default of {@code 10} registered by
+     * {@code OpenSearchSettings.PPL_REX_MAX_MATCH_LIMIT_SETTING}. Cluster-side overrides reach this
+     * map via {@link #setting(String, Object)} — the REST handler reads the live value from {@code
+     * OpenSearchSettings} and routes it through that existing API, keeping {@link
+     * UnifiedQueryContext} decoupled from any specific {@link Settings} implementation.
      */
     private final Map<Settings.Key, Object> settings =
         new HashMap<Settings.Key, Object>(
             Map.of(
                 QUERY_SIZE_LIMIT, SysLimit.DEFAULT.querySizeLimit(),
                 PPL_SUBSEARCH_MAXOUT, SysLimit.DEFAULT.subsearchLimit(),
-                PPL_JOIN_SUBSEARCH_MAXOUT, SysLimit.DEFAULT.joinSubsearchLimit()));
+                PPL_JOIN_SUBSEARCH_MAXOUT, SysLimit.DEFAULT.joinSubsearchLimit(),
+                CALCITE_ENGINE_ENABLED, true,
+                PPL_REX_MAX_MATCH_LIMIT, 10));
 
     /**
      * Sets the query language frontend to be used.

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

@@ -60,7 +60,15 @@ public UnifiedQueryPlanner(UnifiedQueryContext context) {
    */
   public RelNode plan(String query) {
     try {
-      return context.measure(ANALYZE, () -> strategy.plan(query));
+      return context.measure(
+          ANALYZE,
+          () -> {
+            RelNode plan = strategy.plan(query);
+            for (var shuttle : context.getLangSpec().postAnalysisRules()) {
+              plan = plan.accept(shuttle);
+            }
+            return plan;
+          });
     } catch (SyntaxCheckException | UnsupportedOperationException e) {
       throw e;
     } catch (Exception e) {

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

@@ -55,6 +55,11 @@ public PreparedStatement compile(@NonNull RelNode plan) {
   }
 
   private PreparedStatement doCompile(RelNode plan) throws Exception {
+    // Apply pre-compilation rules (e.g., late-binding function impl)
+    for (var rule : context.getLangSpec().preCompilationRules()) {
+      plan = plan.accept(rule);
+    }
+
     // Apply shuttle to convert LogicalTableScan to BindableTableScan
     final RelHomogeneousShuttle shuttle =
         new RelHomogeneousShuttle() {

api/src/main/java/org/opensearch/sql/api/spec/FunctionSpecBuilder.java

@@ -0,0 +1,186 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.api.spec;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.function.BiFunction;
+import javax.annotation.Nullable;
+import lombok.RequiredArgsConstructor;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlFunctionCategory;
+import org.apache.calcite.sql.SqlIdentifier;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperandCountRange;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.parser.SqlParserPos;
+import org.apache.calcite.sql.type.InferTypes;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.SqlOperandCountRanges;
+import org.apache.calcite.sql.type.SqlOperandMetadata;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
+
+/** Fluent DSL for building {@link UnifiedFunctionSpec} instances. */
+@RequiredArgsConstructor
+class FunctionSpecBuilder {
+  /** Function name to register. */
+  private final String name;
+
+  /**
+   * Wraps an existing Calcite operator, preserving its native type system and RexImpTable
+   * implementation for in-memory execution.
+   *
+   * @param op the Calcite operator to delegate to
+   * @return a builder that produces the spec on {@code build()}
+   */
+  DelegateFunctionBuilder delegateTo(SqlOperator op) {
+    return new DelegateFunctionBuilder(name, op);
+  }
+
+  /**
+   * Builds a pushdown-only UDF with relaxed type checking. The resulting function has no local
+   * implementation and delegates execution to the data source via pushdown.
+   *
+   * @param paramNames required parameter names for signature display
+   * @return a builder that produces the spec on {@code build()}
+   */
+  CatalogFunctionBuilder vararg(String... paramNames) {
+    return new CatalogFunctionBuilder(name, List.of(paramNames));
+  }
+
+  /**
+   * Builds a typed SqlFunction with strict operand type checking. Optionally accepts a late-binding
+   * {@code impl} that rewrites the function into executable Calcite expressions at compilation
+   * time.
+   *
+   * @param families operand type families for validation
+   * @return a builder that produces the spec on {@code build()}
+   */
+  DefaultFunctionBuilder operands(SqlTypeFamily... families) {
+    return new DefaultFunctionBuilder(name, families);
+  }
+
+  @RequiredArgsConstructor
+  static class DefaultFunctionBuilder {
+    private final String name;
+    private final SqlTypeFamily[] operandFamilies;
+    private SqlReturnTypeInference returnType;
+    private SqlFunctionCategory category = SqlFunctionCategory.USER_DEFINED_FUNCTION;
+    private @Nullable BiFunction<RexBuilder, RexCall, RexNode> impl;
+
+    DefaultFunctionBuilder returns(SqlReturnTypeInference type) {
+      this.returnType = type;
+      return this;
+    }
+
+    DefaultFunctionBuilder category(SqlFunctionCategory cat) {
+      this.category = cat;
+      return this;
+    }
+
+    /**
+     * Defines how this function executes by rewriting to existing Calcite operators. Applied only
+     * at compilation time (late binding) — the logical plan preserves the original function call.
+     *
+     * @param impl rewrite function that converts this call into executable RexNodes
+     * @return this builder
+     */
+    DefaultFunctionBuilder impl(BiFunction<RexBuilder, RexCall, RexNode> impl) {
+      this.impl = impl;
+      return this;
+    }
+
+    UnifiedFunctionSpec build() {
+      Objects.requireNonNull(returnType, "returns() is required");
+      SqlFunction op =
+          new SqlFunction(
+              name.toUpperCase(),
+              SqlKind.OTHER_FUNCTION,
+              returnType,
+              null,
+              OperandTypes.family(operandFamilies),
+              category);
+      return new UnifiedFunctionSpec(name.toLowerCase(), op, impl);
+    }
+  }
+
+  @RequiredArgsConstructor
+  static class DelegateFunctionBuilder {
+    private final String name;
+    private final SqlOperator operator;
+
+    UnifiedFunctionSpec build() {
+      return new UnifiedFunctionSpec(name.toLowerCase(), operator, null);
+    }
+  }
+
+  @RequiredArgsConstructor
+  static class CatalogFunctionBuilder {
+    private final String name;
+    private final List<String> paramNames;
+    private SqlReturnTypeInference returnType;
+
+    CatalogFunctionBuilder returnType(SqlReturnTypeInference type) {
+      this.returnType = type;
+      return this;
+    }
+
+    UnifiedFunctionSpec build() {
+      Objects.requireNonNull(returnType, "returnType is required");
+      return new UnifiedFunctionSpec(
+          name,
+          new SqlUserDefinedFunction(
+              new SqlIdentifier(name, SqlParserPos.ZERO),
+              SqlKind.OTHER_FUNCTION,
+              returnType,
+              InferTypes.ANY_NULLABLE,
+              new VariadicOperandMetadata(paramNames),
+              List::of), // Pushdown-only: no local implementation
+          null);
+    }
+  }
+
+  /**
+   * Custom operand metadata that bypasses Calcite's built-in type checking. Calcite's {@code
+   * FamilyOperandTypeChecker} rejects variadic calls (CALCITE-5366), so this implementation accepts
+   * any operand types and delegates validation to pushdown.
+   */
+  record VariadicOperandMetadata(List<String> paramNames) implements SqlOperandMetadata {
+
+    @Override
+    public List<String> paramNames() {
+      return paramNames;
+    }
+
+    @Override
+    public List<RelDataType> paramTypes(RelDataTypeFactory tf) {
+      return List.of();
+    }
+
+    @Override
+    public boolean checkOperandTypes(SqlCallBinding binding, boolean throwOnFailure) {
+      return true;
+    }
+
+    @Override
+    public SqlOperandCountRange getOperandCountRange() {
+      return SqlOperandCountRanges.from(paramNames.size());
+    }
+
+    @Override
+    public String getAllowedSignatures(SqlOperator op, String opName) {
+      return opName + "(" + String.join(", ", paramNames) + "[, option=value ...])";
+    }
+  }
+}

api/src/main/java/org/opensearch/sql/api/spec/LanguageSpec.java

@@ -7,6 +7,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import org.apache.calcite.rel.RelShuttle;
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.SqlOperatorTable;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
@@ -17,8 +18,8 @@
 
 /**
  * Language specification defining the dialect the engine accepts. Provides parser configuration,
- * validator configuration, and composable {@link LanguageExtension}s that contribute operators and
- * post-parse rewrite rules.
+ * validator configuration, and composable {@link LanguageExtension}s that contribute operators,
+ * post-parse rewrite rules, and post-analysis rewrite rules.
  *
  * <p>Implementations define a complete language surface — for example, {@link UnifiedSqlSpec}
  * provides ANSI and extended SQL modes. A future PPL spec would implement this same interface once
@@ -27,8 +28,9 @@
 public interface LanguageSpec {
 
   /**
-   * A composable language extension that contributes operators and post-parse rewrite rules. All
-   * methods have defaults so extensions only override what they need.
+   * A composable language extension that contributes operators, post-parse rewrite rules, and
+   * post-analysis rewrite rules. All methods have defaults so extensions only override what they
+   * need.
    */
   interface LanguageExtension {
 
@@ -47,6 +49,23 @@ default SqlOperatorTable operators() {
     default List<SqlVisitor<SqlNode>> postParseRules() {
       return List.of();
     }
+
+    /**
+     * RelNode rewrite rules applied after analysis and before execution. Each rule transforms the
+     * logical plan tree. Rules within a single extension are applied in list order.
+     */
+    default List<RelShuttle> postAnalysisRules() {
+      return List.of();
+    }
+
+    /**
+     * Pre-compilation rules applied only before in-memory execution. Each rule transforms the
+     * logical plan (e.g., binding function implementations). Not applied when the plan is consumed
+     * by external engines.
+     */
+    default List<RelShuttle> preCompilationRules() {
+      return List.of();
+    }
   }
 
   /**
@@ -62,9 +81,9 @@ default List<SqlVisitor<SqlNode>> postParseRules() {
   SqlValidator.Config validatorConfig();
 
   /**
-   * Language extensions registered with this spec. Each extension contributes operators and
-   * post-parse rewrite rules that are composed by {@link #operatorTable()} and {@link
-   * #postParseRules()}.
+   * Language extensions registered with this spec. Each extension contributes operators, post-parse
+   * rewrite rules, and post-analysis rewrite rules composed by {@link #operatorTable()}, {@link
+   * #postParseRules()}, and {@link #postAnalysisRules()}.
    */
   List<LanguageExtension> extensions();
 
@@ -86,4 +105,20 @@ default SqlOperatorTable operatorTable() {
   default List<SqlVisitor<SqlNode>> postParseRules() {
     return extensions().stream().flatMap(ext -> ext.postParseRules().stream()).toList();
   }
+
+  /**
+   * All post-analysis RelNode rewrite rules from registered extensions, flattened in registration
+   * order. Applied to the logical plan after analysis and before execution.
+   */
+  default List<RelShuttle> postAnalysisRules() {
+    return extensions().stream().flatMap(ext -> ext.postAnalysisRules().stream()).toList();
+  }
+
+  /**
+   * All pre-compilation rules from registered extensions, flattened in registration order. Applied
+   * only before in-memory execution.
+   */
+  default List<RelShuttle> preCompilationRules() {
+    return extensions().stream().flatMap(ext -> ext.preCompilationRules().stream()).toList();
+  }
 }