[Analytics-Engine] Seed row counts for tables referenced only inside a subquery

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/IndexRowCountFetcher.java

@@ -10,6 +10,9 @@
 
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.TableScan;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexShuttle;
+import org.apache.calcite.rex.RexSubQuery;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
@@ -64,15 +67,26 @@ private IndexRowCountFetcher() {}
      *     production via {@code DefaultPlanExecutor}'s injected NodeClient.
      */
     public static ToLongFunction<String> fetchFor(RelNode root, Client client) {
-        Set<String> indexNames = new HashSet<>();
-        collectTableScans(root, indexNames);
+        Set<String> indexNames = referencedIndexNames(root);
         if (indexNames.isEmpty() || client == null) {
             return PlannerContext.DEFAULT_TABLE_ROW_COUNTS;
         }
         Map<String, Long> rowCounts = fetchPrimaryDocCounts(client, indexNames);
         return name -> rowCounts.getOrDefault(name, PlannerContext.UNKNOWN_ROW_COUNT);
     }
 
+    /**
+     * Every index name reachable from {@code root}, including those referenced only inside a subquery.
+     *
+     * @param root the RelNode about to be planned
+     * @return the set of index names whose row counts need seeding
+     */
+    public static Set<String> referencedIndexNames(RelNode root) {
+        Set<String> indexNames = new HashSet<>();
+        collectTableScans(root, indexNames);
+        return indexNames;
+    }
+
     private static void collectTableScans(RelNode node, Set<String> out) {
         if (node instanceof TableScan scan) {
             List<String> qualified = scan.getTable().getQualifiedName();
@@ -83,6 +97,19 @@ private static void collectTableScans(RelNode node, Set<String> out) {
         for (RelNode input : node.getInputs()) {
             collectTableScans(input, out);
         }
+        // A subquery is not an input — it hangs off a RexNode (a Filter condition or a Project expression)
+        // as a RexSubQuery, and this fetcher runs BEFORE decorrelation turns it into a join. Walking only
+        // getInputs() therefore misses every table that appears solely inside a subquery, and each of those
+        // scans then falls back to Calcite's default row count. That default is small, so the estimates
+        // derived from it collapse: filters and aggregates above such a scan come out at a row or two, which
+        // makes a large build look tiny and lets plan choices be made on a number with no basis in the data.
+        node.accept(new RexShuttle() {
+            @Override
+            public RexNode visitSubQuery(RexSubQuery subQuery) {
+                collectTableScans(subQuery.rel, out);
+                return super.visitSubQuery(subQuery);
+            }
+        });
     }
 
     private static Map<String, Long> fetchPrimaryDocCounts(Client client, Set<String> indexNames) {

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/IndexRowCountFetcherTests.java

@@ -0,0 +1,97 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.analytics.exec;
+
+import org.apache.calcite.rel.RelNode;
+import org.opensearch.analytics.planner.SqlPlannerTestFixture;
+import org.opensearch.cluster.ClusterState;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Which indices get their row counts seeded.
+ *
+ * <p>This runs BEFORE decorrelation, so a subquery is still a {@code RexSubQuery} hanging off a Filter
+ * condition or a Project expression rather than an input in the RelNode tree. Collecting scans by walking
+ * inputs alone therefore misses any index that appears ONLY inside a subquery, and every such scan then
+ * falls back to Calcite's default row count.
+ *
+ * <p>That default is tiny, so estimates derived from it collapse — a filter and aggregate above such a
+ * scan come out at about one row — and plan choices then get made on a number with no relation to the
+ * data. The failure is silent: the query still returns correct results, just via a plan chosen on
+ * fiction, so nothing but a test like this catches a regression here.
+ */
+public class IndexRowCountFetcherTests extends OpenSearchTestCase {
+
+    private static final Map<String, Map<String, Object>> FIELDS = Map.of(
+        "status",
+        Map.of("type", "integer"),
+        "size",
+        Map.of("type", "integer")
+    );
+
+    private static Set<String> namesFor(String sql, String... indices) {
+        ClusterState state = SqlPlannerTestFixture.clusterStateWith(List.of(indices), FIELDS);
+        RelNode parsed = SqlPlannerTestFixture.parseSql(sql, state);
+        return IndexRowCountFetcher.referencedIndexNames(parsed);
+    }
+
+    /** The plain case: an index scanned by the outer query is collected. */
+    public void testOuterQueryIndexIsCollected() {
+        assertEquals(Set.of("outer_index"), namesFor("SELECT * FROM outer_index", "outer_index"));
+    }
+
+    /**
+     * The regression this class exists for: an index reachable only through an EXISTS subquery must be
+     * collected too. Missing it leaves that scan on Calcite's default row count while its sibling is
+     * correctly sized, so the two sides of a join are estimated on incomparable numbers.
+     */
+    public void testIndexReferencedOnlyInsideAnExistsSubqueryIsCollected() {
+        Set<String> names = namesFor(
+            "SELECT * FROM outer_index WHERE EXISTS (SELECT 1 FROM inner_index WHERE inner_index.status = outer_index.status)",
+            "outer_index",
+            "inner_index"
+        );
+        assertTrue(
+            "a table reachable only through a subquery must still be seeded, else it keeps Calcite's default "
+                + "row count while the outer table is sized correctly; collected "
+                + names,
+            names.contains("inner_index")
+        );
+        assertEquals(Set.of("outer_index", "inner_index"), names);
+    }
+
+    /** Same requirement for an IN subquery, which lowers through a different rule. */
+    public void testIndexReferencedOnlyInsideAnInSubqueryIsCollected() {
+        Set<String> names = namesFor(
+            "SELECT * FROM outer_index WHERE status IN (SELECT status FROM inner_index)",
+            "outer_index",
+            "inner_index"
+        );
+        assertEquals(Set.of("outer_index", "inner_index"), names);
+    }
+
+    /** And for a scalar subquery in the WHERE clause. */
+    public void testIndexReferencedOnlyInsideAScalarSubqueryIsCollected() {
+        Set<String> names = namesFor(
+            "SELECT * FROM outer_index WHERE status > (SELECT max(status) FROM inner_index)",
+            "outer_index",
+            "inner_index"
+        );
+        assertEquals(Set.of("outer_index", "inner_index"), names);
+    }
+
+    /** A subquery over the SAME index must not produce a duplicate or a spurious extra name. */
+    public void testSubqueryOverTheSameIndexCollectsItOnce() {
+        assertEquals(Set.of("outer_index"), namesFor("SELECT * FROM outer_index WHERE EXISTS (SELECT 1 FROM outer_index)", "outer_index"));
+    }
+}

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java

@@ -110,6 +110,43 @@ public static ClusterState clusterStateWith(String indexName, Map<String, Map<St
         return clusterStateWith(indexName, fields, "parquet", 1);
     }
 
+    /**
+     * Builds a {@link ClusterState} holding SEVERAL indices, all with the same fields. Needed whenever a
+     * test has to tell two indices apart — for example a subquery that reads a different index from the
+     * outer query, where a single-index state would make the two indistinguishable.
+     *
+     * @param indexNames the indices to create; all share {@code fields}
+     * @param fields the field mappings applied to every index
+     * @return a cluster state containing all of them, parquet-primary with one shard each
+     */
+    public static ClusterState clusterStateWith(java.util.List<String> indexNames, Map<String, Map<String, Object>> fields) {
+        try {
+            Metadata.Builder metadata = Metadata.builder();
+            for (String indexName : indexNames) {
+                try (XContentBuilder mapping = XContentBuilder.builder(MediaTypeRegistry.JSON.xContent())) {
+                    mapping.startObject().field("properties", fields).endObject();
+                    metadata.put(
+                        IndexMetadata.builder(indexName)
+                            .settings(
+                                Settings.builder()
+                                    .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT.id)
+                                    .put("index.composite.primary_data_format", "parquet")
+                                    .putList("index.composite.secondary_data_formats", "lucene")
+                            )
+                            .numberOfShards(1)
+                            .numberOfReplicas(0)
+                            .putMapping(mapping.toString())
+                            .build(),
+                        false
+                    );
+                }
+            }
+            return ClusterState.builder(new ClusterName("test")).metadata(metadata.build()).build();
+        } catch (Exception e) {
+            throw new AssertionError("Failed to build ClusterState for indices: " + indexNames, e);
+        }
+    }
+
     /**
      * Builds a single-index {@link ClusterState} with explicit primary data format and shard count.
      * The primary format flows through to {@code FieldStorageResolver} and decides which backends