""" TiDB Query Plan comparison logic. Compares two execution plans and determines if they are equivalent. """ import re from dataclasses import dataclass, field from typing import Any, Dict, Optional, Tuple from .plan_types import ( IndexScanOp, Operator, OpType, Plan, PointGetOp, TableScanOp, op_type_is_data_source, ) @dataclass class ComparisonResult: """Result of plan comparison.""" is_same: bool reason: str = "" similarity_score: float = 0.0 # 0.0 to 1.0 est_row_similarity: float = 1.0 # 0.0 to 1.0, how similar are the row estimates details: Dict[str, Any] = field( default_factory=dict ) # Additional comparison details class PlanComparator: """Compares two TiDB execution plans.""" def __init__(self): self.table_aliases: Dict[str, str] = {} def compare_plans(self, plan1: Plan, plan2: Plan) -> ComparisonResult: """Compare two execution plans with both strict and similarity analysis.""" # if plan1.sql != plan2.sql: # return ComparisonResult( # is_same=False, reason="Different SQLs", similarity_score=0.0 # ) # Extract table aliases from SQL self.table_aliases = self._extract_table_aliases(plan1.sql) # Remove projection operators (they don't affect plan semantics) root1 = self._remove_projection(plan1.root) if plan1.root else None root2 = self._remove_projection(plan2.root) if plan2.root else None if not root1 and not root2: return ComparisonResult(is_same=True, similarity_score=1.0) if not root1 or not root2: return ComparisonResult( is_same=False, reason="One plan is empty", similarity_score=0.0 ) # same logic as golang plan-change-capture strict_reason, strict_same = self._compare_operators_strict(root1, root2) # add similarity score similarity_score = self._calculate_operator_similarity(root1, root2) est_row_similarity = self._calculate_est_row_similarity(root1, root2) # calculate the new plan-level estRow similarity plan_estrow_similarity = self.calculate_plan_estrow_similarity(plan1, plan2) # construct detailed result details = { "strict_comparison_result": "passed" if strict_same else "failed", "strict_failure_reason": strict_reason if not strict_same else None, "overall_similarity": similarity_score, "root_est_row_similarity": est_row_similarity, "plan_estrow_similarity": plan_estrow_similarity, # Add new metric "root1_est_rows": root1.get_est_row(), "root2_est_rows": root2.get_est_row(), } return ComparisonResult( is_same=strict_same, reason=strict_reason, similarity_score=similarity_score, est_row_similarity=plan_estrow_similarity, details=details, ) def _remove_projection(self, node: Optional[Operator]) -> Optional[Operator]: """Remove projection operators from the plan tree (compatible with Golang version).""" if not node: return None if node.get_type() == OpType.PROJECTION: children = node.get_children() # Match Golang behavior: directly return first child (assume it exists) return children[0] if children else None # Recursively remove projections from children children = node.get_children() for i, child in enumerate(children): node.set_child(i, self._remove_projection(child)) return node def _compare_operators_strict( self, op1: Operator, op2: Operator ) -> Tuple[str, bool]: # Special handling for PointGet optimization if self._special_handle_point_get(op1, op2): return "", True # Check operator type and task type if op1.get_type() != op2.get_type() or op1.get_task() != op2.get_task(): return ( f"Different operators {op1.get_id()} ({op1.get_type()}) vs {op2.get_id()} ({op2.get_type()})", False, ) # Check children count children1 = op1.get_children() children2 = op2.get_children() if len(children1) != len(children2): return ( f"{op1.get_id()} and {op2.get_id()} have different children count: {len(children1)} vs {len(children2)}", False, ) # Compare operator-specific properties op_reason, op_same = self._compare_operator_specifics_strict(op1, op2) if not op_same: return op_reason, False # Recursively compare children for i in range(len(children1)): # Skip if both children are None if children1[i] is None and children2[i] is None: continue # If only one is None, plans are different if children1[i] is None or children2[i] is None: return f"Child {i} is None in only one plan", False child_reason, child_same = self._compare_operators_strict( children1[i], children2[i] ) if not child_same: return child_reason, False # All comparisons passed return "", True def _compare_operator_specifics_strict( self, op1: Operator, op2: Operator ) -> Tuple[str, bool]: op_type = op1.get_type() if op_type == OpType.TABLE_SCAN: t1, t2 = op1, op2 if isinstance(t1, TableScanOp) and isinstance(t2, TableScanOp): if not self._same_table(t1.table, t2.table): return ( f"Different table scan {t1.get_id()}:{t1.table}, {t2.get_id()}:{t2.table}", False, ) elif op_type == OpType.INDEX_SCAN: i1, i2 = op1, op2 if isinstance(i1, IndexScanOp) and isinstance(i2, IndexScanOp): if not self._same_table(i1.table, i2.table) or i1.index != i2.index: return ( f"Different index scan {i1.get_id()}:{i1.table}:{i1.index}, {i2.get_id()}:{i2.table}:{i2.index}", False, ) return "", True def _calculate_est_row_similarity(self, op1: Operator, op2: Operator) -> float: """Calculate similarity between estimated rows of two operators.""" est1, est2 = op1.get_est_row(), op2.get_est_row() # If both are zero, consider them identical if est1 == 0 and est2 == 0: return 1.0 # If one is zero and the other isn't, they're different if (est1 == 0) != (est2 == 0): return 0.0 # Calculate ratio-based similarity if est1 > 0 and est2 > 0: ratio = min(est1, est2) / max(est1, est2) return ratio return 1.0 def calculate_plan_estrow_similarity(self, plan1: Plan, plan2: Plan) -> float: """ Calculate the overall estRows similarity for the entire plan tree. This method uses a weighted approach that: 1. Gives higher weight to upper-level operators 2. Gives special importance to data source operators 3. Properly handles tree structure differences Returns: float: A similarity score from 0.0 to 1.0 """ if not plan1.root or not plan2.root: return 0.0 if (plan1.root is None) != (plan2.root is None) else 1.0 # Remove projections for consistent comparison root1 = self._remove_projection(plan1.root) root2 = self._remove_projection(plan2.root) total_score, total_weight = self._calculate_weighted_estrow_similarity( root1, root2, level=0 ) # Avoid division by zero if total_weight == 0: return 0.0 return total_score / total_weight def _calculate_weighted_estrow_similarity( self, op1: Operator, op2: Operator, level: int = 0 ) -> tuple: """ Recursively calculate weighted estRow similarity for operators. Args: op1: First operator op2: Second operator level: Current tree level (0 = root) Returns: tuple: (weighted_similarity_score, total_weight) """ if op1 is None or op2 is None: return 0.0, 0.0 # Calculate current node similarity est_similarity = self._calculate_est_row_similarity(op1, op2) # Determine weight based on operator type and level weight = self._get_operator_weight(op1, level) weighted_score = est_similarity * weight # Process children with recursive calls children1 = op1.get_children() children2 = op2.get_children() # Handle different number of children min_children = min(len(children1), len(children2)) children_score = 0.0 children_weight = 0.0 # Compare corresponding children for i in range(min_children): if children1[i] is not None and children2[i] is not None: child_score, child_weight = self._calculate_weighted_estrow_similarity( children1[i], children2[i], level + 1 ) children_score += child_score children_weight += child_weight return weighted_score + children_score, weight + children_weight def _get_operator_weight(self, op: Operator, level: int) -> float: """ Calculate importance weight for an operator. Args: op: The operator level: Tree level (0 = root) Returns: float: Weight value """ # Base weight decreases with depth base_weight = 1.0 * (0.7**level) # Each level reduces weight by 30% # Increase weight for important operator types op_type = op.get_type() # Data source operators are most critical for performance if op_type in {OpType.TABLE_SCAN, OpType.INDEX_SCAN, OpType.POINT_GET}: type_multiplier = 2.0 # Join operators are also very important elif "JOIN" in str(op_type).upper(): type_multiplier = 1.5 # Selection operators often have less reliable estRows elif op_type == OpType.SELECTION: type_multiplier = 0.7 else: type_multiplier = 1.0 # Consider estimated rows (if available) - larger row counts are more impactful est_rows = op.get_est_row() size_factor = min(1.0 + (est_rows / 1000), 2.0) if est_rows > 0 else 1.0 return base_weight * type_multiplier * size_factor def _special_handle_point_get(self, op1: Operator, op2: Operator) -> bool: """Handle special case where PointGet is equivalent to other data sources.""" # PointGet is always better than Table/IndexReader and IndexLookup # when accessing the same table if ( not op_type_is_data_source(op1.get_type()) or op2.get_type() != OpType.POINT_GET ): return False is_simple, table = self._extract_data_source_info(op1) if isinstance(op2, PointGetOp): if is_simple and self._same_table(op2.table, table): return True return False def _extract_data_source_info(self, op: Operator) -> Tuple[bool, str]: """Extract data source information. Returns (is_simple, table_name).""" op_type = op.get_type() if op_type == OpType.TABLE_SCAN: if isinstance(op, TableScanOp): return True, op.table elif op_type == OpType.INDEX_SCAN: if isinstance(op, IndexScanOp): return True, op.table elif op_type in {OpType.TABLE_READER, OpType.INDEX_READER}: children = op.get_children() if children: return self._extract_data_source_info(children[0]) elif op_type == OpType.INDEX_LOOKUP: children = op.get_children() if len(children) >= 2: simple_idx, table1 = self._extract_data_source_info(children[0]) simple_tbl, table2 = self._extract_data_source_info(children[1]) return simple_idx and simple_tbl, table1 return False, "" def _same_table(self, table1: str, table2: str) -> bool: """Check if two table names refer to the same table (considering aliases).""" t1, t2 = table1.lower(), table2.lower() if t1 == t2: return True # Check aliases if self.table_aliases.get(t1) == t2: return True if self.table_aliases.get(t2) == t1: return True return False def _extract_table_aliases(self, sql: str) -> Dict[str, str]: """Extract table aliases from SQL statement.""" aliases = {} # Simple regex-based extraction (can be improved with proper SQL parsing) # Pattern: table_name [AS] alias_name pattern = r"\b(\w+)\s+(?:as\s+)?(\w+)\b" matches = re.findall(pattern, sql.lower(), re.IGNORECASE) for table, alias in matches: # Avoid matching keywords if table not in { "select", "from", "where", "join", "inner", "left", "right", "on", "and", "or", }: aliases[table] = alias return aliases def calculate_similarity(self, plan1: Plan, plan2: Plan) -> float: """Calculate similarity score between two plans (0.0 to 1.0).""" if not plan1.root or not plan2.root: return 0.0 if (plan1.root is None) != (plan2.root is None) else 1.0 return self._calculate_operator_similarity(plan1.root, plan2.root) def _calculate_operator_similarity(self, op1: Operator, op2: Operator) -> float: """calculate similarity score between two operators.""" # Base score for operator type match type_score = 1.0 if op1.get_type() == op2.get_type() else 0.0 task_score = 1.0 if op1.get_task() == op2.get_task() else 0.0 # Children similarity children1 = op1.get_children() children2 = op2.get_children() if len(children1) == len(children2) == 0: children_score = 1.0 elif len(children1) != len(children2): max_children = max(len(children1), len(children2)) min_children = min(len(children1), len(children2)) # calculate child in same position first child_scores = [] for i in range(min_children): # Handle None children if children1[i] is None and children2[i] is None: child_scores.append(1.0) # Same (both None) continue if children1[i] is None or children2[i] is None: child_scores.append(0.0) # Different (only one is None) continue score = self._calculate_operator_similarity(children1[i], children2[i]) child_scores.append(score) # extra child get 0 score for i in range(min_children, max_children): child_scores.append(0.0) children_score = ( sum(child_scores) / len(child_scores) if child_scores else 0.0 ) else: # same children count, iterate each child child_scores = [] for i in range(len(children1)): # Handle None children if children1[i] is None and children2[i] is None: child_scores.append(1.0) # Same (both None) continue if children1[i] is None or children2[i] is None: child_scores.append(0.0) # Different (only one is None) continue score = self._calculate_operator_similarity(children1[i], children2[i]) child_scores.append(score) children_score = ( sum(child_scores) / len(child_scores) if child_scores else 1.0 ) # Operator-specific similarity specific_score = self._calculate_operator_specific_similarity(op1, op2) # EstRow similarity estrow_score = self._calculate_est_row_similarity(op1, op2) # Weighted average (weights are arbitrary, can be adjusted) return ( type_score * 0.25 + task_score * 0.05 + children_score * 0.4 + specific_score * 0.15 + estrow_score * 0.15 ) def _calculate_operator_specific_similarity( self, op1: Operator, op2: Operator ) -> float: """Calculate operator-specific similarity.""" if op1.get_type() != op2.get_type(): return 0.0 op_type = op1.get_type() if op_type == OpType.TABLE_SCAN: if isinstance(op1, TableScanOp) and isinstance(op2, TableScanOp): return 1.0 if self._same_table(op1.table, op2.table) else 0.0 elif op_type == OpType.INDEX_SCAN: if isinstance(op1, IndexScanOp) and isinstance(op2, IndexScanOp): table_match = self._same_table(op1.table, op2.table) index_match = op1.index == op2.index return ( 1.0 if table_match and index_match else (0.5 if table_match else 0.0) ) return 1.0 # Default to full similarity for other operators def compare_plans(plan1: Plan, plan2: Plan) -> ComparisonResult: """Convenient function to compare two plans.""" comparator = PlanComparator() return comparator.compare_plans(plan1, plan2) def calculate_plan_similarity(plan1: Plan, plan2: Plan) -> float: """Convenient function to calculate plan similarity.""" comparator = PlanComparator() return comparator.calculate_similarity(plan1, plan2)