[SNOW-1632895] Add derive_dependent_columns_with_duplication capability

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

@@ -2,11 +2,12 @@
 # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
 #
 
-from typing import AbstractSet, Optional
+from typing import AbstractSet, List, Optional
 
 from snowflake.snowpark._internal.analyzer.expression import (
     Expression,
     derive_dependent_columns,
+    derive_dependent_columns_with_duplication,
 )
 from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
     PlanNodeCategory,
@@ -29,6 +30,9 @@ def __str__(self):
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.left, self.right)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.left, self.right)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.LOW_IMPACT

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

@@ -35,6 +35,13 @@
 def derive_dependent_columns(
     *expressions: "Optional[Expression]",
 ) -> Optional[AbstractSet[str]]:
+    """
+    Given set of expressions, derive the set of columns that the expressions dependents on.
+
+    Note, the returned dependent columns is a set without duplication. For example, given expression
+    concat(col1, upper(co1), upper(col2)), the result will be {col1, col2} even if col1 has
+    occurred in the given expression twice.
+    """
     result = set()
     for exp in expressions:
         if exp is not None:
@@ -48,6 +55,23 @@ def derive_dependent_columns(
     return result
 
 
+def derive_dependent_columns_with_duplication(
+    *expressions: "Optional[Expression]",
+) -> List[str]:
+    """
+    Given set of expressions, derive the list of columns that the expression dependents on.
+
+    Note, the returned columns will have duplication if the column occurred more than once in
+    the given expression. For example, concat(col1, upper(co1), upper(col2)) will have result
+    [col1, col1, col2], where col1 occurred twice in the result.
+    """
+    result = []
+    for exp in expressions:
+        if exp is not None:
+            result.extend(exp.dependent_column_names_with_duplication())
+    return result
+
+
 class Expression:
     """Consider removing attributes, and adding properties and methods.
     A subclass of Expression may have no child, one child, or multiple children.
@@ -68,6 +92,9 @@ def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         # TODO: consider adding it to __init__ or use cached_property.
         return COLUMN_DEPENDENCY_EMPTY
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return []
+
     @property
     def pretty_name(self) -> str:
         """Returns a user-facing string representation of this expression's name.
@@ -143,6 +170,9 @@ def __init__(self, plan: "SnowflakePlan") -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return COLUMN_DEPENDENCY_DOLLAR
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return list(COLUMN_DEPENDENCY_DOLLAR)
+
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
         return self.plan.cumulative_node_complexity
@@ -156,6 +186,9 @@ def __init__(self, expressions: List[Expression]) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(*self.expressions)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(*self.expressions)
+
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
         return sum_node_complexities(
@@ -172,6 +205,9 @@ def __init__(self, columns: Expression, values: List[Expression]) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.columns, *self.values)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.columns, *self.values)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.IN
@@ -212,6 +248,9 @@ def __str__(self):
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return {self.name}
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return [self.name]
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.COLUMN
@@ -235,6 +274,13 @@ def dependent_column_names(self) -> Optional[AbstractSet[str]]:
             else COLUMN_DEPENDENCY_ALL
         )
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return (
+            derive_dependent_columns_with_duplication(*self.expressions)
+            if self.expressions
+            else []  # we currently do not handle * dependency
+        )
+
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
         complexity = {} if self.expressions else {PlanNodeCategory.COLUMN: 1}
@@ -278,6 +324,14 @@ def __hash__(self):
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return self._dependent_column_names
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return (
+            []
+            if (self._dependent_column_names == COLUMN_DEPENDENCY_ALL)
+            or (self._dependent_column_names is None)
+            else list(self._dependent_column_names)
+        )
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.COLUMN
@@ -371,6 +425,9 @@ def __init__(self, expr: Expression, pattern: Expression) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr, self.pattern)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.expr, self.pattern)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         # expr LIKE pattern
@@ -400,6 +457,9 @@ def __init__(
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr, self.pattern)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.expr, self.pattern)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         # expr REG_EXP pattern
@@ -423,6 +483,9 @@ def __init__(self, expr: Expression, collation_spec: str) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.expr)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         # expr COLLATE collate_spec
@@ -444,6 +507,9 @@ def __init__(self, expr: Expression, field: str) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.expr)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         # the literal corresponds to the contribution from self.field
@@ -466,6 +532,9 @@ def __init__(self, expr: Expression, field: int) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.expr)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         # the literal corresponds to the contribution from self.field
@@ -510,6 +579,9 @@ def sql(self) -> str:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(*self.children)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(*self.children)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.FUNCTION
@@ -525,6 +597,9 @@ def __init__(self, expr: Expression, order_by_cols: List[Expression]) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr, *self.order_by_cols)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.expr, *self.order_by_cols)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         # expr WITHIN GROUP (ORDER BY cols)
@@ -549,13 +624,21 @@ def __init__(
         self.branches = branches
         self.else_value = else_value
 
-    def dependent_column_names(self) -> Optional[AbstractSet[str]]:
+    @property
+    def _child_expressions(self) -> List[Expression]:
         exps = []
         for exp_tuple in self.branches:
             exps.extend(exp_tuple)
         if self.else_value is not None:
             exps.append(self.else_value)
-        return derive_dependent_columns(*exps)
+
+        return exps
+
+    def dependent_column_names(self) -> Optional[AbstractSet[str]]:
+        return derive_dependent_columns(*self._child_expressions)
+
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(*self._child_expressions)
 
     @property
     def plan_node_category(self) -> PlanNodeCategory:
@@ -602,6 +685,9 @@ def __init__(
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(*self.children)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(*self.children)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.FUNCTION
@@ -617,6 +703,9 @@ def __init__(self, col: Expression, delimiter: str, is_distinct: bool) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.col)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.col)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.FUNCTION
@@ -636,6 +725,9 @@ def __init__(self, exprs: List[Expression]) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(*self.exprs)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(*self.exprs)
+
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
         return sum_node_complexities(

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

@@ -7,6 +7,7 @@
 from snowflake.snowpark._internal.analyzer.expression import (
     Expression,
     derive_dependent_columns,
+    derive_dependent_columns_with_duplication,
 )
 from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
     PlanNodeCategory,
@@ -23,6 +24,9 @@ def __init__(self, group_by_exprs: List[Expression]) -> None:
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(*self.group_by_exprs)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(*self.group_by_exprs)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.LOW_IMPACT
@@ -45,6 +49,10 @@ def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         flattened_args = [exp for sublist in self.args for exp in sublist]
         return derive_dependent_columns(*flattened_args)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        flattened_args = [exp for sublist in self.args for exp in sublist]
+        return derive_dependent_columns_with_duplication(*flattened_args)
+
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
         return sum_node_complexities(

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

@@ -2,11 +2,12 @@
 # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
 #
 
-from typing import AbstractSet, Optional, Type
+from typing import AbstractSet, List, Optional, Type
 
 from snowflake.snowpark._internal.analyzer.expression import (
     Expression,
     derive_dependent_columns,
+    derive_dependent_columns_with_duplication,
 )
 
 
@@ -55,3 +56,6 @@ def __init__(
 
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.child)
+
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.child)

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

@@ -2,12 +2,13 @@
 # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
 #
 
-from typing import AbstractSet, Dict, Optional
+from typing import AbstractSet, Dict, List, Optional
 
 from snowflake.snowpark._internal.analyzer.expression import (
     Expression,
     NamedExpression,
     derive_dependent_columns,
+    derive_dependent_columns_with_duplication,
 )
 from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
     PlanNodeCategory,
@@ -36,6 +37,9 @@ def __str__(self):
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.child)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.child)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.LOW_IMPACT

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

@@ -7,6 +7,7 @@
 from snowflake.snowpark._internal.analyzer.expression import (
     Expression,
     derive_dependent_columns,
+    derive_dependent_columns_with_duplication,
 )
 from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
     PlanNodeCategory,
@@ -71,6 +72,9 @@ def __init__(
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.lower, self.upper)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.lower, self.upper)
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.LOW_IMPACT
@@ -102,6 +106,11 @@ def dependent_column_names(self) -> Optional[AbstractSet[str]]:
             *self.partition_spec, *self.order_spec, self.frame_spec
         )
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(
+            *self.partition_spec, *self.order_spec, self.frame_spec
+        )
+
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
         # partition_spec order_by_spec frame_spec
@@ -138,6 +147,11 @@ def __init__(
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.window_function, self.window_spec)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(
+            self.window_function, self.window_spec
+        )
+
     @property
     def plan_node_category(self) -> PlanNodeCategory:
         return PlanNodeCategory.WINDOW
@@ -171,6 +185,9 @@ def __init__(
     def dependent_column_names(self) -> Optional[AbstractSet[str]]:
         return derive_dependent_columns(self.expr, self.default)
 
+    def dependent_column_names_with_duplication(self) -> List[str]:
+        return derive_dependent_columns_with_duplication(self.expr, self.default)
+
     @property
     def individual_node_complexity(self) -> Dict[PlanNodeCategory, int]:
         # for func_name