feat: enable Calcite engine for Prometheus datasources with filter pushdown

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -251,18 +251,32 @@ public RelNode visitRelation(Relation node, CalcitePlanContext context) {
     DataSourceSchemaIdentifierNameResolver nameResolver =
         new DataSourceSchemaIdentifierNameResolver(
             dataSourceService, node.getTableQualifiedName().getParts());
-    if (!nameResolver
-        .getDataSourceName()
-        .equals(DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME)) {
-      throw new CalciteUnsupportedException(
-          "Datasource " + nameResolver.getDataSourceName() + " is unsupported in Calcite");
-    }
     if (nameResolver.getIdentifierName().equals(DATASOURCES_TABLE_NAME)) {
       throw new CalciteUnsupportedException("SHOW DATASOURCES is unsupported in Calcite");
     }
     if (nameResolver.getSchemaName().equals(INFORMATION_SCHEMA_NAME)) {
       throw new CalciteUnsupportedException("information_schema is unsupported in Calcite");
     }
+    // For non-default datasources, verify the table supports Calcite integration
+    // before proceeding. If it doesn't, fall back to V2 via CalciteUnsupportedException.
+    if (!nameResolver
+        .getDataSourceName()
+        .equals(DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME)) {
+      org.opensearch.sql.storage.Table storageTable =
+          dataSourceService
+              .getDataSource(nameResolver.getDataSourceName())
+              .getStorageEngine()
+              .getTable(
+                  new org.opensearch.sql.DataSourceSchemaName(
+                      nameResolver.getDataSourceName(), nameResolver.getSchemaName()),
+                  nameResolver.getIdentifierName());
+      if (!(storageTable instanceof org.apache.calcite.schema.Table)) {
+        throw new CalciteUnsupportedException(
+            "Datasource "
+                + nameResolver.getDataSourceName()
+                + " is unsupported in Calcite (table does not implement Calcite Table interface)");
+      }
+    }
     context.relBuilder.scan(node.getTableQualifiedName().getParts());
     RelNode scan = context.relBuilder.peek();
 

core/src/main/java/org/opensearch/sql/calcite/OpenSearchSchema.java

@@ -9,6 +9,7 @@
 import java.util.Map;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
+import org.apache.calcite.schema.Schema;
 import org.apache.calcite.schema.Table;
 import org.apache.calcite.schema.impl.AbstractSchema;
 import org.opensearch.sql.DataSourceSchemaName;
@@ -34,6 +35,20 @@ public Table get(Object key) {
         }
       };
 
+  private final Map<String, Schema> subSchemaMap =
+      new HashMap<>() {
+        @Override
+        public Schema get(Object key) {
+          if (!super.containsKey(key)) {
+            String dsName = (String) key;
+            if (dataSourceService.dataSourceExists(dsName)) {
+              super.put(dsName, new DataSourceSubSchema(dataSourceService, dsName));
+            }
+          }
+          return super.get(key);
+        }
+      };
+
   public void registerTable(QualifiedName qualifiedName) {
     DataSourceSchemaIdentifierNameResolver nameResolver =
         new DataSourceSchemaIdentifierNameResolver(dataSourceService, qualifiedName.getParts());
@@ -45,6 +60,64 @@ public void registerTable(QualifiedName qualifiedName) {
                 new DataSourceSchemaName(
                     nameResolver.getDataSourceName(), nameResolver.getSchemaName()),
                 nameResolver.getIdentifierName());
-    tableMap.put(qualifiedName.toString(), (org.apache.calcite.schema.Table) table);
+    if (table instanceof org.apache.calcite.schema.Table calciteTable) {
+      tableMap.put(qualifiedName.toString(), calciteTable);
+    } else {
+      throw new UnsupportedOperationException(
+          "Table "
+              + qualifiedName
+              + " does not support Calcite integration. "
+              + "The storage engine table must implement org.apache.calcite.schema.Table.");
+    }
+  }
+
+  /**
+   * A sub-schema representing a non-default datasource. Lazily resolves tables from the
+   * datasource's storage engine, allowing Calcite to find tables via schema-qualified names like
+   * scan(["prometheus", "up"]).
+   */
+  private static class DataSourceSubSchema extends AbstractSchema {
+    private final DataSourceService dataSourceService;
+    private final String dataSourceName;
+
+    DataSourceSubSchema(DataSourceService dataSourceService, String dataSourceName) {
+      this.dataSourceService = dataSourceService;
+      this.dataSourceName = dataSourceName;
+    }
+
+    @Override
+    protected Map<String, Table> getTableMap() {
+      return tableMap;
+    }
+
+    private final Map<String, Table> tableMap =
+        new HashMap<>() {
+          @Override
+          public Table get(Object key) {
+            if (!super.containsKey(key)) {
+              resolveTable((String) key);
+            }
+            return super.get(key);
+          }
+        };
+
+    private void resolveTable(String tableName) {
+      org.opensearch.sql.storage.Table table =
+          dataSourceService
+              .getDataSource(dataSourceName)
+              .getStorageEngine()
+              .getTable(new DataSourceSchemaName(dataSourceName, "default"), tableName);
+      if (table instanceof org.apache.calcite.schema.Table calciteTable) {
+        tableMap.put(tableName, calciteTable);
+      } else {
+        throw new UnsupportedOperationException(
+            "Table "
+                + dataSourceName
+                + "."
+                + tableName
+                + " does not support Calcite integration. "
+                + "The storage engine table must implement org.apache.calcite.schema.Table.");
+      }
+    }
   }
 }

prometheus/build.gradle

@@ -18,6 +18,7 @@ dependencies {
     api project(':core')
     implementation project(':datasources')
     implementation project(':direct-query-core')
+    compileOnly 'org.immutables:value-annotations:2.8.8'
     testImplementation(testFixtures(project(":direct-query-core")))
 
     implementation group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}"
@@ -66,7 +67,11 @@ jacocoTestCoverageVerification {
             element = 'CLASS'
             excludes = [
                     'org.opensearch.sql.prometheus.data.constants.*',
-                    'org.opensearch.sql.prometheus.functions.implementation.*'
+                    'org.opensearch.sql.prometheus.functions.implementation.*',
+                    'org.opensearch.sql.prometheus.planner.logical.rules.EnumerablePrometheusScanRule',
+                    'org.opensearch.sql.prometheus.planner.logical.rules.PrometheusFilterPushDownRule',
+                    'org.opensearch.sql.prometheus.planner.logical.rules.PrometheusRules',
+                    'org.opensearch.sql.prometheus.storage.scan.*'
             ]
             limit {
                 counter = 'LINE'

prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/EnumerablePrometheusScanRule.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.prometheus.planner.logical.rules;
+
+import org.apache.calcite.adapter.enumerable.EnumerableConvention;
+import org.apache.calcite.plan.Convention;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.convert.ConverterRule;
+import org.opensearch.sql.prometheus.storage.scan.CalciteEnumerablePrometheusScan;
+import org.opensearch.sql.prometheus.storage.scan.CalciteLogicalPrometheusScan;
+
+/**
+ * Rule to convert a {@link CalciteLogicalPrometheusScan} to a {@link
+ * CalciteEnumerablePrometheusScan}.
+ */
+public class EnumerablePrometheusScanRule extends ConverterRule {
+
+  /** Default configuration. */
+  public static final Config DEFAULT_CONFIG =
+      Config.INSTANCE
+          .as(Config.class)
+          .withConversion(
+              CalciteLogicalPrometheusScan.class,
+              s -> s.getPrometheusTable() != null,
+              Convention.NONE,
+              EnumerableConvention.INSTANCE,
+              "EnumerablePrometheusScanRule")
+          .withRuleFactory(EnumerablePrometheusScanRule::new);
+
+  /** Creates an EnumerablePrometheusScanRule. */
+  protected EnumerablePrometheusScanRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public boolean matches(RelOptRuleCall call) {
+    return true;
+  }
+
+  @Override
+  public RelNode convert(RelNode rel) {
+    final CalciteLogicalPrometheusScan scan = (CalciteLogicalPrometheusScan) rel;
+    return new CalciteEnumerablePrometheusScan(
+        scan.getCluster(),
+        scan.getTraitSet().plus(EnumerableConvention.INSTANCE),
+        scan.getTable(),
+        scan.getPrometheusTable(),
+        scan.getSchema(),
+        scan.getPushDownContext());
+  }
+}

prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusFilterPushDownRule.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.prometheus.planner.logical.rules;
+
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelOptRuleOperand;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.logical.LogicalFilter;
+import org.opensearch.sql.calcite.utils.PlanUtils;
+import org.opensearch.sql.prometheus.storage.scan.CalciteLogicalPrometheusScan;
+
+/**
+ * Planner rule that pushes filter conditions (time range and label matchers) down into a {@link
+ * CalciteLogicalPrometheusScan}.
+ *
+ * <p>Supported pushdowns:
+ *
+ * <ul>
+ *   <li>Time range comparisons on @timestamp (>, >=, <, <=)
+ *   <li>Label equality conditions (label = 'value')
+ * </ul>
+ *
+ * <p>Unsupported conditions remain as a LogicalFilter on top.
+ */
+public class PrometheusFilterPushDownRule extends RelOptRule {
+
+  public static final PrometheusFilterPushDownRule INSTANCE =
+      new PrometheusFilterPushDownRule(
+          operand(
+              LogicalFilter.class,
+              operand(CalciteLogicalPrometheusScan.class, none())),
+          "PrometheusFilterPushDownRule");
+
+  private PrometheusFilterPushDownRule(RelOptRuleOperand operand, String description) {
+    super(operand, description);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final LogicalFilter filter = call.rel(0);
+    final CalciteLogicalPrometheusScan scan = call.rel(1);
+
+    RelNode newNode = scan.pushDownFilter(filter);
+    if (newNode != null) {
+      call.transformTo(newNode);
+      PlanUtils.tryPruneRelNodes(call);
+    }
+  }
+}