Add query planner Phase 2: logical planning and cost estimation

server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java

@@ -39,6 +39,7 @@ protected FeatureFlagSettings(
         FeatureFlags.TERM_VERSION_PRECOMMIT_ENABLE_SETTING,
         FeatureFlags.ARROW_STREAMS_SETTING,
         FeatureFlags.STREAM_TRANSPORT_SETTING,
-        FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING
+        FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING,
+        FeatureFlags.QUERY_PLANNER_SETTING
     );
 }

server/src/main/java/org/opensearch/common/util/FeatureFlags.java

@@ -120,6 +120,12 @@ public class FeatureFlags {
     public static final String ARROW_STREAMS = FEATURE_FLAG_PREFIX + "arrow.streams.enabled";
     public static final Setting<Boolean> ARROW_STREAMS_SETTING = Setting.boolSetting(ARROW_STREAMS, false, Property.NodeScope);
 
+    /**
+     * Gates the query planner functionality.
+     */
+    public static final String QUERY_PLANNER = FEATURE_FLAG_PREFIX + "query_planner.enabled";
+    public static final Setting<Boolean> QUERY_PLANNER_SETTING = Setting.boolSetting(QUERY_PLANNER, false, Property.NodeScope);
+
     /**
      * Underlying implementation for feature flags.
      * All settable feature flags are tracked here in FeatureFlagsImpl.featureFlags.
@@ -146,6 +152,7 @@ static class FeatureFlagsImpl {
                 put(ARROW_STREAMS_SETTING, ARROW_STREAMS_SETTING.getDefault(Settings.EMPTY));
                 put(STREAM_TRANSPORT_SETTING, STREAM_TRANSPORT_SETTING.getDefault(Settings.EMPTY));
                 put(MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING, MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING.getDefault(Settings.EMPTY));
+                put(QUERY_PLANNER_SETTING, QUERY_PLANNER_SETTING.getDefault(Settings.EMPTY));
             }
         };
 

server/src/main/java/org/opensearch/search/planner/AbstractQueryPlanNode.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.search.planner;
+
+import org.apache.lucene.search.Query;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Base implementation of QueryPlanNode providing common functionality.
+ *
+ * @opensearch.internal
+ */
+public abstract class AbstractQueryPlanNode implements QueryPlanNode {
+
+    protected final Query luceneQuery;
+    protected final QueryNodeType nodeType;
+    protected final List<QueryPlanNode> children;
+    protected volatile QueryCost estimatedCost;
+
+    protected AbstractQueryPlanNode(Query luceneQuery, QueryNodeType nodeType) {
+        this(luceneQuery, nodeType, Collections.emptyList());
+    }
+
+    protected AbstractQueryPlanNode(Query luceneQuery, QueryNodeType nodeType, List<QueryPlanNode> children) {
+        this.luceneQuery = luceneQuery;
+        this.nodeType = nodeType;
+        this.children = Collections.unmodifiableList(children);
+    }
+
+    @Override
+    public Query getLuceneQuery() {
+        return luceneQuery;
+    }
+
+    @Override
+    public QueryNodeType getType() {
+        return nodeType;
+    }
+
+    @Override
+    public List<QueryPlanNode> getChildren() {
+        return children;
+    }
+
+    @Override
+    public QueryCost getEstimatedCost() {
+        QueryCost cost = estimatedCost;
+        if (cost == null) {
+            cost = calculateCost();
+            estimatedCost = cost;
+        }
+        return cost;
+    }
+
+    /**
+     * Calculate the cost for this node. Subclasses should implement
+     * their specific cost calculation logic.
+     */
+    protected abstract QueryCost calculateCost();
+
+    @Override
+    public QueryPlanProfile getProfile() {
+        QueryPlanProfile profile = new QueryPlanProfile(nodeType.name(), getEstimatedCost());
+        addProfileAttributes(profile);
+        return profile;
+    }
+
+    /**
+     * Add node-specific attributes to the profile.
+     * Subclasses can override to add custom attributes.
+     */
+    protected void addProfileAttributes(QueryPlanProfile profile) {
+        profile.addAttribute("lucene_query", luceneQuery.toString());
+        profile.addAttribute("children_count", children.size());
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append(nodeType.name()).append("[");
+        sb.append("query=").append(luceneQuery.getClass().getSimpleName());
+        sb.append(", cost=").append(getEstimatedCost().getTotalCost());
+        if (!children.isEmpty()) {
+            sb.append(", children=").append(children.size());
+        }
+        sb.append("]");
+        return sb.toString();
+    }
+}

server/src/main/java/org/opensearch/search/planner/CostEstimator.java

@@ -0,0 +1,195 @@
+/*
+ * 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.search.planner;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.Weight;
+import org.opensearch.index.query.QueryShardContext;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Estimates query costs using Lucene's scorer costs and OpenSearch heuristics.
+ * @opensearch.internal
+ */
+public class CostEstimator {
+
+    private static final Logger logger = LogManager.getLogger(CostEstimator.class);
+
+    private final QueryShardContext queryShardContext;
+    private final IndexSearcher searcher;
+    private final IndexReader reader;
+
+    // Cost multipliers for different query types
+    private static final Map<String, CostMultipliers> QUERY_COST_MULTIPLIERS = new HashMap<>();
+
+    // Fallback cost fractions when actual costs can't be computed
+    private static final double FALLBACK_TERM_FRACTION = 0.008;  // ~0.8% of docs for medium selectivity term
+    private static final double FALLBACK_RANGE_FRACTION = 0.12;   // Ranges typically match more docs
+    private static final double FALLBACK_WILDCARD_FRACTION = 0.18; // Wildcards are expensive
+    private static final double FALLBACK_FUZZY_FRACTION = 0.35;    // Fuzzy matches are very expensive
+    private static final double FALLBACK_DEFAULT_FRACTION = 0.05;  // Conservative default
+
+    static {
+        // Initialize cost multipliers based on query complexity
+        QUERY_COST_MULTIPLIERS.put("Script", new CostMultipliers(12.5, 2.3, 1.0));
+        QUERY_COST_MULTIPLIERS.put("Fuzzy", new CostMultipliers(4.8, 3.2, 1.1));
+        QUERY_COST_MULTIPLIERS.put("Regexp", new CostMultipliers(5.2, 2.8, 1.0));
+        QUERY_COST_MULTIPLIERS.put("Wildcard", new CostMultipliers(2.3, 1.0, 1.8));
+        QUERY_COST_MULTIPLIERS.put("Prefix", new CostMultipliers(1.7, 1.0, 1.6));
+        QUERY_COST_MULTIPLIERS.put("Range", new CostMultipliers(1.0, 1.0, 1.4));
+        QUERY_COST_MULTIPLIERS.put("Nested", new CostMultipliers(2.1, 2.2, 1.0));
+        QUERY_COST_MULTIPLIERS.put("FunctionScore", new CostMultipliers(3.5, 1.2, 1.0));
+    }
+
+    private static class CostMultipliers {
+        final double cpu;
+        final double memory;
+        final double io;
+
+        CostMultipliers(double cpu, double memory, double io) {
+            this.cpu = cpu;
+            this.memory = memory;
+            this.io = io;
+        }
+    }
+
+    public CostEstimator(QueryShardContext queryShardContext) {
+        this.queryShardContext = queryShardContext;
+        this.searcher = queryShardContext.searcher();
+        this.reader = searcher != null ? searcher.getIndexReader() : null;
+    }
+
+    /**
+     * Estimates the document frequency for a term.
+     */
+    public long estimateTermDocFrequency(String field, String value) {
+        if (reader == null) {
+            // Default estimate if no reader available
+            logger.trace("No reader available for term frequency estimation, using default");
+            return 100;
+        }
+
+        try {
+            Term term = new Term(field, value);
+            long docFreq = reader.docFreq(term);
+            logger.trace("Term frequency for field '{}' value '{}': {}", field, value, docFreq);
+            return docFreq;
+        } catch (IOException e) {
+            logger.debug("Error estimating term frequency for field '{}': {}", field, e.getMessage());
+            // Fall back to estimate
+            return 100;
+        }
+    }
+
+    /**
+     * Gets the total number of documents in the index.
+     */
+    public long getTotalDocs() {
+        if (reader == null) {
+            return 10000; // Default estimate
+        }
+        return reader.maxDoc();
+    }
+
+    /**
+     * Estimates cost for a generic query using Lucene's Weight API.
+     */
+    public long estimateGenericCost(Query query) {
+        if (searcher == null) {
+            // Fallback estimate based on query type
+            logger.trace("No searcher available, using fallback cost estimation");
+            return estimateFallbackCost(query);
+        }
+
+        try {
+            // Create weight to get scorer and cost
+            Weight weight = searcher.createWeight(searcher.rewrite(query), org.apache.lucene.search.ScoreMode.COMPLETE, 1);
+
+            // Aggregate costs across all leaves
+            if (reader != null && !reader.leaves().isEmpty()) {
+                long totalCost = 0;
+                int leafCount = 0;
+                for (LeafReaderContext leafContext : reader.leaves()) {
+                    Scorer scorer = weight.scorer(leafContext);
+                    if (scorer != null) {
+                        long leafCost = scorer.iterator().cost();
+                        totalCost += leafCost;
+                        leafCount++;
+                    }
+                }
+                if (totalCost > 0) {
+                    logger.trace("Lucene cost for {} across {} leaves: {}", query.getClass().getSimpleName(), leafCount, totalCost);
+                    return totalCost;
+                }
+            }
+        } catch (IOException e) {
+            logger.debug("Error estimating Lucene cost for query {}: {}", query.getClass().getSimpleName(), e.getMessage());
+            // Fall back to estimate
+        }
+
+        return estimateFallbackCost(query);
+    }
+
+    private long estimateFallbackCost(Query query) {
+        // Heuristic estimates when we can't get actual costs
+        String queryType = query.getClass().getSimpleName();
+
+        // Base estimates as percentage of total docs
+        long totalDocs = getTotalDocs();
+
+        if (queryType.contains("MatchAllDocs")) {
+            return totalDocs;
+        } else if (queryType.contains("Term")) {
+            return (long) (totalDocs * FALLBACK_TERM_FRACTION);
+        } else if (queryType.contains("Range")) {
+            return (long) (totalDocs * FALLBACK_RANGE_FRACTION);
+        } else if (queryType.contains("Wildcard") || queryType.contains("Prefix")) {
+            return (long) (totalDocs * FALLBACK_WILDCARD_FRACTION);
+        } else if (queryType.contains("Fuzzy") || queryType.contains("Regexp")) {
+            return (long) (totalDocs * FALLBACK_FUZZY_FRACTION);
+        }
+
+        return (long) (totalDocs * FALLBACK_DEFAULT_FRACTION);
+    }
+
+    public QueryCost estimateFullCost(Query query, long luceneCost) {
+        String queryType = query.getClass().getSimpleName();
+
+        // Find matching cost multipliers
+        CostMultipliers multipliers = null;
+        for (Map.Entry<String, CostMultipliers> entry : QUERY_COST_MULTIPLIERS.entrySet()) {
+            if (queryType.contains(entry.getKey())) {
+                multipliers = entry.getValue();
+                break;
+            }
+        }
+
+        // Use defaults if no specific multipliers found
+        double cpuMultiplier = multipliers != null ? multipliers.cpu : 1.0;
+        double memoryMultiplier = multipliers != null ? multipliers.memory : 1.0;
+        double ioMultiplier = multipliers != null ? multipliers.io : 1.0;
+
+        // Base costs
+        double cpuCost = 0.001 * cpuMultiplier;
+        long memoryCost = (long) (luceneCost * 8 * memoryMultiplier);
+        double ioCost = 0.01 * ioMultiplier;
+
+        return new QueryCost(luceneCost, cpuCost, memoryCost, ioCost, true);
+    }
+}