CNDB-15508: Address review comments

Author: Piotr Kołaczkowski

src/java/org/apache/cassandra/config/CassandraRelevantProperties.java

@@ -629,6 +629,13 @@ public enum CassandraRelevantProperties
      */
     SAI_QUERY_KIND_PER_QUERY_METRICS_ENABLED("cassandra.sai.metrics.query_kind.per_query.enabled", "false"),
 
+    /**
+     * Whether to enable SAI query plan metrics such as the estimated cost, estimated number of rows,
+     * number of indexes used in the original and optimized query plan, etc.
+     * These metrics are counters and histograms.
+     */
+    SAI_QUERY_PLAN_METRICS_ENABLED("cassandra.sai.metrics.query_plan.enabled", "true"),
+
     /**
      * Whether to enable SAI index metrics such as memtable flush metrics, compaction metrics, and disk usage metrics.
      * These metrics include timers, histograms, counters, and gauges for index operations.

src/java/org/apache/cassandra/index/sai/QueryContext.java

@@ -26,6 +26,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 
+import org.apache.cassandra.config.CassandraRelevantProperties;
 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.index.sai.plan.Plan;
 import org.apache.cassandra.index.sai.utils.AbortedOperationException;
@@ -69,10 +70,7 @@ public class QueryContext
 
     private final LongAdder shadowedPrimaryKeyCount = new LongAdder();
 
-    private Plan originalPlan = null;
-
-
-    private Plan optimizedPlan = null;
+    private PlanInfo queryPlanInfo;
 
     @VisibleForTesting
     public QueryContext()
@@ -207,18 +205,10 @@ public long annGraphSearchLatency()
         return annGraphSearchLatency.longValue();
     }
 
-    /** Can return null when query plan hasn't been prepared and optimized yet */
-    @Nullable
-    public Plan optimizedPlan()
-    {
-        return optimizedPlan;
-    }
-
-    /** Can return null when query plan hasn't been prepared yet */
     @Nullable
-    public Plan originalPlan()
+    public PlanInfo queryPlanInfo()
     {
-        return originalPlan;
+        return queryPlanInfo;
     }
 
     public void checkpoint()
@@ -254,14 +244,10 @@ public void updateAnnRerankFloor(float observedFloor)
             annRerankFloor = max(annRerankFloor, observedFloor);
     }
 
-    public void setOriginalPlan(Plan originalPlan)
+    public void recordQueryPlan(Plan.RowsIteration originalPlan, Plan.RowsIteration.RowsIteration optimizedPlan)
     {
-        this.originalPlan = originalPlan;
-    }
-
-    public void setOptimizedPlan(Plan optimizedPlan)
-    {
-        this.optimizedPlan = optimizedPlan;
+        if (CassandraRelevantProperties.SAI_QUERY_PLAN_METRICS_ENABLED.getBoolean())
+            this.queryPlanInfo = new PlanInfo(originalPlan, optimizedPlan);
     }
 
     public Snapshot snapshot()
@@ -297,7 +283,7 @@ public static class Snapshot
         public final long shadowedPrimaryKeyCount;
 
         @Nullable
-        public final QueryPlanInfo queryPlanInfo;
+        public final PlanInfo queryPlanInfo;
 
         /**
          * Creates a snapshot of all the metrics in the given {@link QueryContext}.
@@ -322,40 +308,34 @@ private Snapshot(QueryContext context)
             queryTimeouts = context.queryTimeouts();
             annGraphSearchLatency = context.annGraphSearchLatency();
             shadowedPrimaryKeyCount = context.getShadowedPrimaryKeyCount();
-
-            Plan originalPlan = context.originalPlan();
-            Plan optimizedPlan = context.optimizedPlan();
-            if (originalPlan != null && optimizedPlan != null) {
-                queryPlanInfo = new QueryPlanInfo(originalPlan, optimizedPlan);
-            } else {
-                queryPlanInfo = null;
-            }
+            queryPlanInfo = context.queryPlanInfo();
         }
+    }
+
+    /**
+     * Captures relevant information about a query plan, both original and optimized.
+     */
+    public static class PlanInfo
+    {
+        public final boolean searchExecutedBeforeOrder;
+        public final boolean filterExecutedAfterOrderedScan;
 
-        public static class QueryPlanInfo
+        public final double rowsEstimated;
+        public final double selectivityEstimated;
+        public final double costEstimated;
+
+        public final int indexReferencesInQuery;
+        public final int indexReferencesInPlan;
+
+        public PlanInfo(@Nonnull Plan.RowsIteration originalPlan, @Nonnull Plan.RowsIteration optimizedPlan)
         {
-            public final boolean searchExecutedBeforeOrder;
-            public final boolean filterExecutedAfterOrderedScan;
-
-            public final double rowsEstimated;
-            public final double selectivityEstimated;
-            public final double costEstimated;
-
-            public final int indexReferencesInQuery;
-            public final int indexReferencesInPlan;
-
-            public QueryPlanInfo(@Nonnull Plan originalPlan, @Nonnull Plan optimizedPlan)
-            {
-                this.costEstimated = optimizedPlan.fullCost();
-                this.rowsEstimated = optimizedPlan.expectedRows();
-                this.selectivityEstimated = optimizedPlan.selectivity();
-                this.indexReferencesInQuery = originalPlan.referencedIndexCount();
-                this.indexReferencesInPlan = optimizedPlan.referencedIndexCount();
-                this.searchExecutedBeforeOrder = optimizedPlan.isSearchThenOrderHybrid();
-                this.filterExecutedAfterOrderedScan = optimizedPlan.isOrderedScanThenFilterHybrid();
-            }
+            this.costEstimated = optimizedPlan.fullCost();
+            this.rowsEstimated = optimizedPlan.expectedRows();
+            this.selectivityEstimated = optimizedPlan.selectivity();
+            this.indexReferencesInQuery = originalPlan.referencedIndexCount();
+            this.indexReferencesInPlan = optimizedPlan.referencedIndexCount();
+            this.searchExecutedBeforeOrder = optimizedPlan.isSearchThenOrderHybrid();
+            this.filterExecutedAfterOrderedScan = optimizedPlan.isOrderedScanThenFilterHybrid();
         }
-
     }
-
 }

src/java/org/apache/cassandra/index/sai/metrics/TableQueryMetrics.java

@@ -233,7 +233,7 @@ public void record(QueryContext.Snapshot snapshot)
             totalPartitionReads.inc(snapshot.partitionsRead);
             totalRowsFiltered.inc(snapshot.rowsFiltered);
 
-            QueryContext.Snapshot.QueryPlanInfo queryPlanInfo = snapshot.queryPlanInfo;
+            QueryContext.PlanInfo queryPlanInfo = snapshot.queryPlanInfo;
             if (queryPlanInfo != null)
             {
                 totalCostEstimated.inc(Math.round(queryPlanInfo.costEstimated));
@@ -381,7 +381,7 @@ public void record(QueryContext.Snapshot snapshot)
 
             shadowedKeysScannedHistogram.update(snapshot.shadowedPrimaryKeyCount);
 
-            QueryContext.Snapshot.QueryPlanInfo queryPlanInfo = snapshot.queryPlanInfo;
+            QueryContext.PlanInfo queryPlanInfo = snapshot.queryPlanInfo;
             if (queryPlanInfo != null)
             {
                 costEstimated.update(Math.round(queryPlanInfo.costEstimated));

src/java/org/apache/cassandra/index/sai/plan/Plan.java

@@ -36,7 +36,6 @@
 import org.apache.cassandra.db.filter.IndexHints;
 import org.apache.cassandra.db.filter.RowFilter;
 import org.apache.cassandra.index.sai.IndexContext;
-import org.apache.cassandra.index.sai.QueryContext;
 import org.apache.cassandra.index.sai.disk.format.Version;
 import org.apache.cassandra.index.sai.iterators.KeyRangeIntersectionIterator;
 import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
@@ -395,7 +394,7 @@ protected void visitIndexes(Consumer<IndexContext> consumer)
      * and recomputes the nodes above it. Then it returns the best plan from candidates obtained that way.
      * The expected running time is proportional to the height of the plan tree multiplied by the number of the leaves.
      */
-    public final Plan optimize()
+    protected Plan optimize()
     {
         if (logger.isTraceEnabled())
             logger.trace("Optimizing plan:\n{}", toRedactedStringRecursive());
@@ -428,7 +427,7 @@ public final Plan optimize()
      * Modifies all intersections to not intersect more clauses than the given limit.
      * Retains the most selective clauses.
      */
-    public final Plan limitIntersectedClauses(int clauseLimit)
+    protected Plan limitIntersectedClauses(int clauseLimit)
     {
         Plan result = this;
         if (result instanceof Intersection)
@@ -529,19 +528,6 @@ public final double selectivity()
         return selectivity;
     }
 
-    /**
-     * Returns the number of rows produced by this plan.
-     * This can be only called on plans producing rows.
-     */
-    public final double expectedRows()
-    {
-        Cost cost = cost();
-        if (!(cost instanceof RowsIterationCost))
-            throw new UnsupportedOperationException("Expected rows is only supported for plans returning rows (called on " + this.getClass() + ')');
-
-        return ((RowsIterationCost) cost).expectedRows;
-    }
-
     /**
      * Returns the number of indexes referenced by this plan.
      */
@@ -729,6 +715,18 @@ final double costPerKey()
             return cost().costPerKey();
         }
 
+        @Override
+        public final KeysIteration optimize()
+        {
+            return (KeysIteration) super.optimize();
+        }
+
+        @Override
+        public final KeysIteration limitIntersectedClauses(int clauseLimit)
+        {
+            return (KeysIteration) super.limitIntersectedClauses(clauseLimit);
+        }
+
         protected abstract boolean usesIncludedIndex();
     }
 
@@ -1639,6 +1637,26 @@ final double costPerRow()
         {
             return cost().costPerRow();
         }
+
+        /**
+         * Returns the number of rows produced by this plan.
+         */
+        public final double expectedRows()
+        {
+            return cost().expectedRows;
+        }
+
+        @Override
+        public final RowsIteration optimize()
+        {
+            return (RowsIteration) super.optimize();
+        }
+
+        @Override
+        public final RowsIteration limitIntersectedClauses(int clauseLimit)
+        {
+            return (RowsIteration) super.limitIntersectedClauses(clauseLimit);
+        }
     }
 
     /**

src/java/org/apache/cassandra/index/sai/plan/QueryController.java

@@ -394,14 +394,14 @@ Plan buildPlan()
         // in which predicates it leaves in the plan and the probability of accidentally removing a good branch
         // here is even lower.
         int intersectionClauseLimit = CassandraRelevantProperties.SAI_INTERSECTION_CLAUSE_LIMIT.getInt();
-        Plan plan = rowsIteration.limitIntersectedClauses(intersectionClauseLimit * 3);
-        queryContext.setOriginalPlan(plan);
+        Plan.RowsIteration origPlan = rowsIteration.limitIntersectedClauses(intersectionClauseLimit * 3);
+        Plan.RowsIteration plan = origPlan;
 
         if (QUERY_OPT_LEVEL > 0)
-            plan = plan.optimize();
+            plan = origPlan.optimize();
 
         plan = plan.limitIntersectedClauses(intersectionClauseLimit);
-        queryContext.setOptimizedPlan(plan);
+        queryContext.recordQueryPlan(origPlan, plan);
         updateIndexMetricsQueriesCount(plan);
 
         if (logger.isTraceEnabled())