diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/IndexRowCountFetcher.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/IndexRowCountFetcher.java index 5adc21f4456fe..af1fadc7a70ae 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/IndexRowCountFetcher.java +++ b/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,8 +67,7 @@ private IndexRowCountFetcher() {} * production via {@code DefaultPlanExecutor}'s injected NodeClient. */ public static ToLongFunction fetchFor(RelNode root, Client client) { - Set indexNames = new HashSet<>(); - collectTableScans(root, indexNames); + Set indexNames = referencedIndexNames(root); if (indexNames.isEmpty() || client == null) { return PlannerContext.DEFAULT_TABLE_ROW_COUNTS; } @@ -73,6 +75,18 @@ public static ToLongFunction fetchFor(RelNode root, Client client) { 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 referencedIndexNames(RelNode root) { + Set indexNames = new HashSet<>(); + collectTableScans(root, indexNames); + return indexNames; + } + private static void collectTableScans(RelNode node, Set out) { if (node instanceof TableScan scan) { List qualified = scan.getTable().getQualifiedName(); @@ -83,6 +97,19 @@ private static void collectTableScans(RelNode node, Set 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 fetchPrimaryDocCounts(Client client, Set indexNames) { diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/IndexRowCountFetcherTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/IndexRowCountFetcherTests.java new file mode 100644 index 0000000000000..13ab5f76ef2ff --- /dev/null +++ b/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. + * + *

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. + * + *

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> FIELDS = Map.of( + "status", + Map.of("type", "integer"), + "size", + Map.of("type", "integer") + ); + + private static Set 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 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 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 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")); + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java index 5aaf239786c40..704d015c2754d 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java @@ -110,6 +110,43 @@ public static ClusterState clusterStateWith(String indexName, Map indexNames, Map> 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