SNOW-1432019 Calculate subtree query complexity (#1657)

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -939,7 +939,7 @@ def do_resolve_with_resolved_children(
                 )
 
         if isinstance(logical_plan, UnresolvedRelation):
-            return self.plan_builder.table(logical_plan.name)
+            return self.plan_builder.table(logical_plan.name, logical_plan)
 
         if isinstance(logical_plan, SnowflakeCreateTable):
             return self.plan_builder.save_as_table(

src/snowflake/snowpark/_internal/analyzer/binary_expression.py

@@ -8,6 +8,9 @@
     Expression,
     derive_dependent_columns,
 )
+from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
+    PlanNodeCategory,
+)
 
 
 class BinaryExpression(Expression):
@@ -26,6 +29,10 @@ def __str__(self):
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.left, self.right)
 
+    @property
+    def plan_node_category(self) -> PlanNodeCategory:
+        return PlanNodeCategory.LOW_IMPACT
+
 
 class BinaryArithmeticExpression(BinaryExpression):
     pass

src/snowflake/snowpark/_internal/analyzer/binary_plan_node.py

@@ -2,9 +2,13 @@
 # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
 #
 
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from snowflake.snowpark._internal.analyzer.expression import Expression
+from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
+    PlanNodeCategory,
+    sum_node_complexities,
+)
 from snowflake.snowpark._internal.analyzer.snowflake_plan_node import LogicalPlan
 from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 
@@ -69,7 +73,10 @@ def __init__(self, left: LogicalPlan, right: LogicalPlan) -> None:
 
 
 class SetOperation(BinaryNode):
-    pass
+    @property
+    def plan_node_category(self) -> PlanNodeCategory:
+        # (left) operator (right)
+        return PlanNodeCategory.SET_OPERATION
 
 
 class Except(SetOperation):
@@ -187,3 +194,33 @@ def __init__(
     @property
     def sql(self) -> str:
         return self.join_type.sql
+
+    @property
+    def plan_node_category(self) -> PlanNodeCategory:
+        return PlanNodeCategory.JOIN
+
+    @property
+    def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
+        # SELECT * FROM (left) AS left_alias join_type_sql JOIN (right) AS right_alias match_cond, using_cond, join_cond
+        complexity = {self.plan_node_category: 1}
+        if isinstance(self.join_type, UsingJoin) and self.join_type.using_columns:
+            complexity = sum_node_complexities(
+                complexity,
+                {PlanNodeCategory.COLUMN: len(self.join_type.using_columns)},
+            )
+        complexity = (
+            sum_node_complexities(
+                complexity, self.join_condition.cumulative_node_complexity
+            )
+            if self.join_condition
+            else complexity
+        )
+
+        complexity = (
+            sum_node_complexities(
+                complexity, self.match_condition.cumulative_node_complexity
+            )
+            if self.match_condition
+            else complexity
+        )
+        return complexity