snowpark-connect 0.27.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/snowpark_connect/column_name_handler.py

@@ -13,18 +13,21 @@ from functools import cached_property
 from pyspark.errors.exceptions.base import AnalysisException
 
 from snowflake.snowpark import DataFrame
-from snowflake.snowpark._internal.analyzer.analyzer_utils import (
-    quote_name_without_upper_casing,
-    unquote_if_quoted,
-)
+from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark._internal.utils import quote_name
 from snowflake.snowpark.types import StructType
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 from snowflake.snowpark_connect.config import global_config
-from snowflake.snowpark_connect.hidden_column import HiddenColumn
-from snowflake.snowpark_connect.utils.context import get_current_operation_scope
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.utils.context import (
+    get_current_operation_scope,
+    get_is_processing_order_by,
+)
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
+from snowflake.snowpark_connect.utils.sequence import next_unique_num
 
 ALREADY_QUOTED = re.compile('^(".+")$', re.DOTALL)
 
@@ -44,6 +47,7 @@ def set_schema_getter(df: DataFrame, get_schema: Callable[[], StructType]) -> No
     df.__class__ = PatchedDataFrame
 
 
+# TODO replace plan_id-offset with single unique value
 def make_column_names_snowpark_compatible(
     names: list[str], plan_id: int, offset: int = 0
 ) -> list[str]:
@@ -76,42 +80,42 @@ def make_column_names_snowpark_compatible(
         In this case the function call should be `make_column_names_snowpark_compatible(['a', 'b'], 5, 2)`,
         to avoid naming conflicts between the new columns and the old columns.
     """
+    from snowflake.snowpark_connect.relation.read.metadata_utils import (
+        METADATA_FILENAME_COLUMN,
+    )
+
     return [
+        # Skip METADATA$FILENAME - preserve original name without quoting
+        name if name == METADATA_FILENAME_COLUMN else
         # Use `-` in the name to force df.column to return double-quoted names
         quote_name(f"{unquote_if_quoted(name)}-{plan_id:08x}-{i + offset}")
         for i, name in enumerate(names)
     ]
 
 
+def make_unique_snowpark_name(spark_name: str) -> str:
+    """
+    Returns a snowpark column name that's guaranteed to be unique in this session,
+    by appending "#<unique number>" to the given spark name.
+    """
+    return quote_name(f"{spark_name}-{next_unique_num():x}")
+
+
 @dataclass(frozen=True)
 class ColumnNames:
     spark_name: str
     snowpark_name: str
-    qualifiers: list[str]
+    qualifiers: set[ColumnQualifier]
+    equivalent_snowpark_names: set[str] | None = ((None,),)
     catalog_info: str | None = None  # Catalog from fully qualified name
     database_info: str | None = None  # Database from fully qualified name
+    is_hidden: bool = False  # Hidden columns are only accessible via qualified names
 
-
-def get_list_of_spark_names_for_column(column_names: ColumnNames) -> list[str]:
-    """
-    Returns a list of Spark names for a given ColumnNames object.
-    This is useful when a single Spark name maps to multiple names due to table alias.
-
-    For example, if the column name is 'id' and the qualifiers are ['db', 'table'],
-    then the possible Spark names are:
-      ['id', 'db.table.id', 'table.id']
-    """
-    spark_name = column_names.spark_name
-    qualifiers = column_names.qualifiers
-
-    qualifier_suffixes_list = [
-        ".".join(quote_name_without_upper_casing(x) for x in qualifiers[i:])
-        for i in range(len(qualifiers))
-    ]
-    return [spark_name] + [
-        f"{qualifier_suffix}.{spark_name}"
-        for qualifier_suffix in qualifier_suffixes_list
-    ]
+    def all_spark_names_including_qualified_names(self):
+        all_names = [self.spark_name]
+        for qualifier in self.qualifiers:
+            all_names.extend(qualifier.all_qualified_names(self.spark_name))
+        return all_names
 
 
 class ColumnNameMap:
@@ -123,32 +127,32 @@ class ColumnNameMap:
             [], bool
         ] = lambda: global_config.spark_sql_caseSensitive,
         column_metadata: dict | None = None,
-        column_qualifiers: list[list[str]] | None = None,
-        hidden_columns: set[HiddenColumn] | None = None,
+        column_qualifiers: list[set[ColumnQualifier]] = None,
         parent_column_name_map: ColumnNameMap | None = None,
+        equivalent_snowpark_names: list[set[str]] | None = None,
+        column_is_hidden: list[bool] | None = None,
     ) -> None:
         """
         spark_column_names: Original spark column names
         snowpark_column_names: Snowpark column names
-        column_metadata: This field is used to store metadata related to columns. Since Snowpark’s Struct type does not support metadata,
+        column_metadata: This field is used to store metadata related to columns. Since Snowpark's Struct type does not support metadata,
         we use this attribute to store any metadata related to the columns.
         The key is the original Spark column name, and the value is the metadata.
         example: Dict('age', {'foo': 'bar'})
         column_qualifiers: Optional qualifiers for the columns, used to handle table aliases or DataFrame aliases.
-        hidden_columns: Optional set of HiddenColumn objects.
         parent_column_name_map: parent ColumnNameMap
+        column_is_hidden: Optional list of booleans indicating whether each column is hidden
         """
         self.columns: list[ColumnNames] = []
-        self.spark_to_col = defaultdict(list)
+        self.spark_to_col: defaultdict[str, list[ColumnNames]] = defaultdict(list)
         self.uppercase_spark_to_col = defaultdict(list)
         self.snowpark_to_col = defaultdict(list)
         self.is_case_sensitive = is_case_sensitive
         self.column_metadata = column_metadata
-        self.hidden_columns = hidden_columns
 
         # Rename chain dictionary to track column renaming history
         self.rename_chains: dict[str, str] = {}  # old_name -> new_name mapping
-        self.current_columns: set[str] = set()  # Current column names
+        self.current_columns: set[str] = set()  # current column names
 
         # Parent ColumnNameMap classes
         self._parent_column_name_map = parent_column_name_map
@@ -179,21 +183,22 @@ class ColumnNameMap:
             c = ColumnNames(
                 spark_name=spark_name,
                 snowpark_name=snowpark_column_names[i],
-                qualifiers=column_qualifiers[i] if column_qualifiers else [],
+                qualifiers=column_qualifiers[i]
+                if column_qualifiers and column_qualifiers[i]
+                else set(),
+                equivalent_snowpark_names=equivalent_snowpark_names[i]
+                if equivalent_snowpark_names and equivalent_snowpark_names[i]
+                else set(),
                 catalog_info=catalog_info,
                 database_info=database_info,
+                is_hidden=column_is_hidden[i] if column_is_hidden else False,
             )
             self.columns.append(c)
 
-            # we want to store all the spark names including qualifiers (these are generated from table alias or dataframe alias)
-            spark_names_including_qualifier = get_list_of_spark_names_for_column(c)
-
-            for spark_name_including_qualifier in spark_names_including_qualifier:
+            for spark_name in c.all_spark_names_including_qualified_names():
                 # the same spark name can map to multiple snowpark names
-                self.spark_to_col[spark_name_including_qualifier].append(c)
-                self.uppercase_spark_to_col[
-                    spark_name_including_qualifier.upper()
-                ].append(c)
+                self.spark_to_col[spark_name].append(c)
+                self.uppercase_spark_to_col[spark_name.upper()].append(c)
 
             # the same snowpark name can map to multiple spark column
             # e.g. df.select(date_format('dt', 'yyy'), date_format('dt', 'yyyy')) ->
@@ -286,9 +291,10 @@ class ColumnNameMap:
         self,
         spark_column_names: list[str],
         return_first: bool = False,
+        original_snowpark_names: list[str] | None = None,
     ) -> list[str]:
         snowpark_column_names = self._get_snowpark_column_names_from_spark_column_names(
-            spark_column_names, return_first
+            spark_column_names, return_first, original_snowpark_names
         )
         if snowpark_column_names:
             return snowpark_column_names
@@ -302,7 +308,7 @@ class ColumnNameMap:
             and self._parent_column_name_map is not None
         ):
             snowpark_column_names = self._parent_column_name_map.get_snowpark_column_names_from_spark_column_names(
-                spark_column_names, return_first
+                spark_column_names, return_first, original_snowpark_names
             )
 
         return snowpark_column_names
@@ -311,9 +317,10 @@ class ColumnNameMap:
         self,
         spark_column_names: list[str],
         return_first: bool = False,
+        original_snowpark_names: list[str] | None = None,
     ) -> list[str]:
         snowpark_column_names = []
-        for name in spark_column_names:
+        for i, name in enumerate(spark_column_names):
             if not global_config.spark_sql_caseSensitive:
                 name = name.upper()
                 mapping = self.uppercase_spark_to_col
@@ -325,8 +332,26 @@ class ColumnNameMap:
 
             columns = mapping[name]
 
+            # make sure the column matches the original snowpark name, if given
+            if original_snowpark_names:
+                oname = original_snowpark_names[i]
+                columns = [
+                    c
+                    for c in columns
+                    if c.snowpark_name == oname or oname in c.equivalent_snowpark_names
+                ]
+
+            # Filter out hidden columns for unqualified lookups
+            # A qualified lookup contains a dot (e.g., "b.id"), unqualified doesn't (e.g., "id")
+            # Hidden columns should only be accessible via qualified names
+            is_qualified_lookup = "." in name or original_snowpark_names
+            if not is_qualified_lookup:
+                # Unqualified lookup: only include visible columns
+                columns = [c for c in columns if not c.is_hidden]
+
             if return_first:
-                snowpark_column_names.append(columns[0].snowpark_name)
+                if columns:  # Only append if we have columns after filtering
+                    snowpark_column_names.append(columns[0].snowpark_name)
             else:
                 snowpark_column_names.extend([c.snowpark_name for c in columns])
 
@@ -338,8 +363,7 @@ class ColumnNameMap:
         *,
         allow_non_exists: bool = False,
         return_first: bool = False,
-        is_qualified: bool = False,
-        source_qualifiers: list[str] | None = None,
+        original_snowpark_name: str | None = None,
     ) -> str | None:
         assert isinstance(spark_column_name, str)
         resolved_name = (
@@ -347,52 +371,85 @@ class ColumnNameMap:
             if self.rename_chains
             else spark_column_name
         )
-
-        # We need to check hidden columns first. We want to avoid the code path
-        # within get_snowpark_column_names_from_spark_column_names that checks the parent ColumnNameMap.
-        # This is because that will return the name of the using column that's been dropped from the result
-        # dataframe. We want to fetch and resolve the hidden column to its visible using column name instead.
-        # Even if this is an unqualified reference or one to the visible column, it will resolve correctly to
-        # the visible name anyway.
-        snowpark_names = []
-        # Only check hidden columns for qualified references with source qualifiers
-        if is_qualified and source_qualifiers is not None and self.hidden_columns:
-            column_name = spark_column_name
-
-            # Check each hidden column for column name AND qualifier match
-            for hidden_col in self.hidden_columns:
-                if (
-                    hidden_col.spark_name == column_name
-                    and hidden_col.qualifiers == source_qualifiers
-                ):
-                    if not global_config.spark_sql_caseSensitive:
-                        if hidden_col.spark_name.upper() == column_name.upper() and [
-                            q.upper() for q in hidden_col.qualifiers
-                        ] == [q.upper() for q in source_qualifiers]:
-                            snowpark_names.append(hidden_col.visible_snowpark_name)
-                    else:
-                        snowpark_names.append(hidden_col.visible_snowpark_name)
-
-        # If not found in hidden columns, proceed with normal lookup
-        if not snowpark_names:
-            snowpark_names = self.get_snowpark_column_names_from_spark_column_names(
-                [resolved_name], return_first
-            )
+        snowpark_names = self.get_snowpark_column_names_from_spark_column_names(
+            [resolved_name],
+            return_first,
+            [original_snowpark_name] if original_snowpark_name else None,
+        )
 
         snowpark_names_len = len(snowpark_names)
         if snowpark_names_len > 1:
-            raise AnalysisException(
-                f"Ambiguous spark column name {spark_column_name}, potential snowpark column names {snowpark_names}"
-            )
+            # Check if this is a case where we have identical expressions that can be safely resolved to the first one
+            # This commonly happens with GROUP BY expressions that also appear in SELECT clauses
+            if (
+                get_is_processing_order_by()
+                and self._can_resolve_ambiguous_identical_expressions(
+                    resolved_name, snowpark_names
+                )
+            ):
+                # All the ambiguous columns represent the same expression, so we can safely use the first one
+                return snowpark_names[0]
+            else:
+                exception = AnalysisException(
+                    f"Ambiguous spark column name {spark_column_name}, potential snowpark column names {snowpark_names}"
+                )
+                attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+                raise exception
         elif snowpark_names_len == 0:
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Spark column name {spark_column_name} does not exist"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
         return snowpark_names[0]
 
+    def _can_resolve_ambiguous_identical_expressions(
+        self, spark_column_name: str, snowpark_names: list[str]
+    ) -> bool:
+        """
+        Determine if ambiguous columns represent identical expressions that can be safely resolved to the first one.
+
+        This handles the common case where the same expression (like a UDF call) appears multiple times
+        in a SELECT clause within a GROUP BY query. Since they're the same expression operating on the
+        same grouped data, they will have identical values, so we can safely resolve to any of them.
+
+        Args:
+            spark_column_name: The Spark column name that has multiple mappings, make sure resolve this reforehand
+            snowpark_names: List of Snowpark column names that map to this Spark column name
+
+        Returns:
+            True if we can safely resolve to the first snowpark column, False otherwise
+        """
+        if spark_column_name not in self.spark_to_col:
+            return False
+
+        columns: list[ColumnNames] = self.spark_to_col[spark_column_name]
+
+        # If we don't have multiple columns, there's no ambiguity to resolve
+        if len(columns) <= 1:
+            return False
+
+        # Check if all the snowpark names correspond to columns that have identical underlying expressions
+        # We'll compare the actual column objects to see if they represent the same computation
+        first_column = columns[0]
+
+        for column in columns[1:]:
+            if first_column.qualifiers != column.qualifiers:
+                return False
+
+        # Additional safety check: ensure all snowpark names are actually in our mapping
+        for snowpark_name in snowpark_names:
+            if snowpark_name not in self.snowpark_to_col:
+                return False
+
+        # If we reach here, the columns appear to be identical expressions from the same context
+        # This commonly happens in GROUP BY scenarios where the same expression appears in both
+        # the grouping clause and the select clause
+        return True
+
     def get_spark_column_names_from_snowpark_column_names(
         self,
         snowpark_column_names: list[str],
@@ -418,98 +475,79 @@ class ColumnNameMap:
         )
         spark_names_len = len(spark_names)
         if spark_names_len > 1:
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"Ambiguous snowpark column name {snowpark_column_name}, potential spark column names {spark_names}"
             )
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
         elif spark_names_len == 0:
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Snowpark column name {snowpark_column_name} does not exist"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
         return spark_names[0]
 
     def get_spark_column_name(self, idx: int) -> str:
         return self.columns[idx].spark_name
 
     def get_spark_columns(self) -> list[str]:
-        return [c.spark_name for c in self.columns]
+        return [c.spark_name for c in self.columns if not c.is_hidden]
 
     def get_spark_and_snowpark_columns_with_qualifier_for_qualifier(
-        self, qualifiers_input: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
+        self, target_qualifier: ColumnQualifier
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]]]:
         """
-        Returns the Spark and Snowpark column names along with their qualifiers for the specified qualifiers.
-        If a column does not have a qualifier, it will be None.
+        Returns the Spark and Snowpark column names along with their qualifiers for the specified qualifier.
         """
-        spark_columns = []
-        snowpark_columns = []
-        qualifiers = []
+        spark_columns: list[str] = []
+        snowpark_columns: list[str] = []
+        qualifiers: list[set[ColumnQualifier]] = []
 
+        normalized_qualifier = target_qualifier
         if not self.is_case_sensitive():
-            qualifiers_input = [q.upper() for q in qualifiers_input]
+            normalized_qualifier = target_qualifier.to_upper()
 
-        for c in self.columns:
-            col_qualifiers = (
-                [q.upper() for q in c.qualifiers]
+        for column in self.columns:
+            # Normalize all qualifiers for comparison
+            column_qualifiers: set[ColumnQualifier] = (
+                {q.to_upper() for q in iter(column.qualifiers)}
                 if not self.is_case_sensitive()
-                else c.qualifiers
+                else column.qualifiers
             )
-            if len(col_qualifiers) < len(qualifiers_input):
-                # If the column has fewer qualifiers than the input, it cannot match
-                continue
-            if col_qualifiers[-len(qualifiers_input) :] == qualifiers_input:
-                spark_columns.append(c.spark_name)
-                snowpark_columns.append(c.snowpark_name)
-                qualifiers.append(c.qualifiers)
-
-        # Note: The following code is commented out because there is a bug with handling duplicate columns in
-        # qualified select *'s. This needs to be revisited once a solution for that is found.
-        # TODO: https://snowflakecomputing.atlassian.net/browse/SNOW-2265240
-
-        # # Handles fetching/resolving the hidden columns if they also match the qualifiers
-        # # This method is only ever called for qualified references, so we need to check hidden columns as well.
-        # if self.hidden_columns:
-        #     for hidden_col in self.hidden_columns:
-        #         col_qualifiers = (
-        #             [q.upper() for q in hidden_col.qualifiers]
-        #             if not self.is_case_sensitive()
-        #             else hidden_col.qualifiers
-        #         )
-        #         if len(col_qualifiers) < len(qualifiers_input):
-        #             continue
-        #         if col_qualifiers[-len(qualifiers_input) :] == qualifiers_input:
-        #             # This hidden column matches! Add it to the results
-        #             spark_columns.append(hidden_col.spark_name)
-        #             snowpark_columns.append(hidden_col.visible_snowpark_name)
-        #             qualifiers.append(hidden_col.qualifiers)
+            if any([q.matches(normalized_qualifier) for q in column_qualifiers]):
+                spark_columns.append(column.spark_name)
+                snowpark_columns.append(column.snowpark_name)
+                qualifiers.append(column.qualifiers)
 
         return spark_columns, snowpark_columns, qualifiers
 
     def get_snowpark_columns(self) -> list[str]:
-        return [c.snowpark_name for c in self.columns]
+        return [c.snowpark_name for c in self.columns if not c.is_hidden]
 
-    def get_snowpark_columns_after_drop(self, cols_to_drop: list[str]) -> list[str]:
+    def get_snowpark_columns_after_drop(
+        self, cols_to_drop: list[str]
+    ) -> list[ColumnNames]:
         return [
             c
-            for c in self.get_snowpark_columns()
-            if self._quote_if_unquoted(c) not in cols_to_drop
+            for c in self.columns
+            if self._quote_if_unquoted(c.snowpark_name) not in cols_to_drop
         ]
 
-    def get_qualifiers(self) -> list[list[str]]:
+    def get_qualifiers(self) -> list[set[ColumnQualifier]]:
         """
         Returns the qualifiers for the columns.
-        If a column does not have a qualifier, it will be None.
         """
-        return [c.qualifiers for c in self.columns]
+        return [c.qualifiers for c in self.columns if not c.is_hidden]
 
     def get_qualifiers_for_columns_after_drop(
         self, cols_to_drop: list[str]
-    ) -> list[list[str]]:
+    ) -> list[set[ColumnQualifier]]:
         """
         Returns the qualifiers for the columns after dropping the specified columns.
-        If a column is dropped, its qualifier will be None.
         """
         return [
             c.qualifiers
@@ -517,27 +555,40 @@ class ColumnNameMap:
             if self._quote_if_unquoted(c.snowpark_name) not in cols_to_drop
         ]
 
-    def get_qualifier_for_spark_column(
+    def get_qualifiers_for_snowpark_column(
         self,
-        spark_column_name: str,
-    ) -> list[str]:
+        snowpark_name: str,
+    ) -> set[ColumnQualifier]:
         """
-        Returns the qualifier for the specified Spark column name.
-        If the column does not exist, returns None.
+        Returns the qualifier for the specified snowpark column name.
+        If the column does not exist, returns empty ColumnQualifier.
         """
-        if not self.is_case_sensitive():
-            name = spark_column_name.upper()
-            mapping = self.uppercase_spark_to_col
-        else:
-            name = spark_column_name
-            mapping = self.spark_to_col
+        for c in self.columns:
+            if c.snowpark_name == snowpark_name:
+                return c.qualifiers
+
+        return set()
+
+    def get_equivalent_snowpark_names(self) -> list[set[str]]:
+        return [c.equivalent_snowpark_names for c in self.columns]
 
-        col = mapping.get(name)
+    def get_equivalent_snowpark_names_for_snowpark_name(
+        self, snowpark_name: str | None
+    ) -> set[str]:
+        """
+        Helper method to get the set of old, equivalent snowpark names for the given column. Used to pass
+        this information to child column maps.
+        """
+        if not snowpark_name:
+            return set()
 
-        if col is None or len(col) == 0:
-            return []
+        name = self._quote_if_unquoted(snowpark_name)
+        for c in self.columns:
+            if name == c.snowpark_name:
+                return c.equivalent_snowpark_names
 
-        return col[0].qualifiers
+        # no equivalent names found
+        return set()
 
     @staticmethod
     def _quote_if_unquoted(s: str) -> str:
@@ -555,19 +606,20 @@ class ColumnNameMap:
     def snowpark_to_spark_map(self) -> dict[str, str]:
         return {c.snowpark_name: c.spark_name for c in self.columns}
 
-    def spark_to_snowpark_for_pattern(self, pattern: str) -> list[tuple[str, str]]:
-        pattern_regex = re.compile(
-            pattern, 0 if self.is_case_sensitive() else re.IGNORECASE
-        )
-        return [
-            (c.spark_name, c.snowpark_name)
-            for c in self.columns
-            if pattern_regex.fullmatch(c.spark_name)
-        ]
+    def get_columns_matching_pattern(self, pattern: str) -> list[ColumnNames]:
+        try:
+            pattern_regex = re.compile(
+                pattern, 0 if self.is_case_sensitive() else re.IGNORECASE
+            )
+            return [c for c in self.columns if pattern_regex.fullmatch(c.spark_name)]
+        except re.error as e:
+            exception = AnalysisException(f"Invalid regex pattern '{pattern}': {e}")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT)
+            raise exception
 
     def with_columns(
         self, new_spark_columns: list[str], new_snowpark_columns: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]], list[set[str]]]:
         """
         Returns an ordered list of spark and snowpark column names after adding the new columns through a withColumns call.
         All replaced columns retain their ordering in the dataframe. The new columns are added to the end of the list.
@@ -588,6 +640,7 @@ class ColumnNameMap:
         snowpark_columns = []
         removed_index: set[int] = set()
         qualifiers = []
+        equivalent_snowpark_names = []
 
         for c in self.columns:
             column_name = self._normalized_spark_name(c.spark_name)
@@ -596,19 +649,22 @@ class ColumnNameMap:
                 removed_index.add(index)
                 spark_columns.append(new_spark_columns[index])
                 snowpark_columns.append(new_snowpark_columns[index])
-                qualifiers.append([])
+                qualifiers.append(set())
+                equivalent_snowpark_names.append(set())
             else:
                 spark_columns.append(c.spark_name)
                 snowpark_columns.append(c.snowpark_name)
                 qualifiers.append(c.qualifiers)
+                equivalent_snowpark_names.append(c.equivalent_snowpark_names)
 
         for i, _ in enumerate(new_spark_columns):
             if i not in removed_index:
                 spark_columns.append(new_spark_columns[i])
                 snowpark_columns.append(new_snowpark_columns[i])
-                qualifiers.append([])
+                qualifiers.append(set())
+                equivalent_snowpark_names.append(set())
 
-        return spark_columns, snowpark_columns, qualifiers
+        return spark_columns, snowpark_columns, qualifiers, equivalent_snowpark_names
 
     def _normalized_spark_name(self, spark_name: str) -> str:
         if self.is_case_sensitive():
@@ -616,34 +672,77 @@ class ColumnNameMap:
         else:
             return spark_name.upper()
 
-    def is_hidden_column_reference(
-        self, spark_column_name: str, source_qualifiers: list[str] | None = None
-    ) -> bool:
+    def get_columns_after_join(
+        self, right: ColumnNameMap, join_columns: list[str], join_type: str
+    ) -> list[ColumnNames]:
         """
-        Check if a column reference would be resolved through hidden columns.
+        Returns a list of columns (names and qualifiers) after a using_columns join with the given column map
         """
-        if not self.hidden_columns or source_qualifiers is None:
-            return False
 
-        # For qualified references with source_qualifiers
-        column_name = (
-            spark_column_name  # When has_plan_id=True, this is just the column name
-        )
+        # first, let's gather right-side join columns for qualifier lookup
+        # and the remaining columns to append them to the result
+        join_column_names = [self._normalized_spark_name(c) for c in join_columns]
+        right_join_columns: dict[str, ColumnNames] = {}
+        right_remaining_columns: list[ColumnNames] = []
+        for oc in right.columns:
+            col_name = self._normalized_spark_name(oc.spark_name)
+            # only take the first matching column
+            if col_name in join_column_names and col_name not in right_join_columns:
+                right_join_columns[col_name] = oc
+            else:
+                right_remaining_columns.append(oc)
 
-        for hidden_col in self.hidden_columns:
-            if (
-                hidden_col.spark_name == column_name
-                and hidden_col.qualifiers == source_qualifiers
-            ):
-                if not global_config.spark_sql_caseSensitive:
-                    if hidden_col.spark_name.upper() == column_name.upper() and [
-                        q.upper() for q in hidden_col.qualifiers
-                    ] == [q.upper() for q in source_qualifiers]:
-                        return True
-                else:
-                    return True
+        # now gather left-side columns
+        left_join_columns: dict[str, ColumnNames] = {}
+        left_remaining_columns: list[ColumnNames] = []
+        for c in self.columns:
+            col_name = self._normalized_spark_name(c.spark_name)
+            if col_name in join_column_names and col_name not in left_join_columns:
+                equivalent_snowpark_names = set()
+                # only assign join-side qualifier for outer joins
+                match join_type:
+                    case "left":
+                        qualifiers = c.qualifiers
+                    case "right":
+                        qualifiers = right_join_columns[col_name].qualifiers
+                    case _:
+                        qualifiers = (
+                            c.qualifiers | right_join_columns[col_name].qualifiers
+                        )
+                        equivalent_snowpark_names.update(
+                            c.equivalent_snowpark_names,
+                            right_join_columns[col_name].equivalent_snowpark_names,
+                            {right_join_columns[col_name].snowpark_name},
+                        )
+
+                left_join_columns[col_name] = ColumnNames(
+                    c.spark_name, c.snowpark_name, qualifiers, equivalent_snowpark_names
+                )
+            else:
+                left_remaining_columns.append(c)
+
+        # join columns go first in the user-given order,
+        # then the remaining left-side columns, then remaining right-side columns
+        match join_type:
+            case "right":
+                ordered_join_columns = [
+                    right_join_columns[name] for name in join_column_names
+                ]
+            case _:
+                ordered_join_columns = [
+                    left_join_columns[name] for name in join_column_names
+                ]
+        return ordered_join_columns + left_remaining_columns + right_remaining_columns
 
-        return False
+    def get_conflicting_snowpark_columns(self, other: ColumnNameMap) -> set[str]:
+        conflicting_columns = set()
+        snowpark_names = {c.snowpark_name for c in self.columns}
+
+        for c in other.columns:
+            if c.snowpark_name in snowpark_names:
+                conflicting_columns.add(c.snowpark_name)
+
+        return conflicting_columns
 
 
 class JoinColumnNameMap(ColumnNameMap):
@@ -654,9 +753,6 @@ class JoinColumnNameMap(ColumnNameMap):
     ) -> None:
         self.left_column_mapping: ColumnNameMap = left_colmap
         self.right_column_mapping: ColumnNameMap = right_colmap
-        # Ensure attributes expected by base-class helpers exist to avoid AttributeError
-        # when generic code paths (e.g., hidden column checks) touch them.
-        self.hidden_columns: set[HiddenColumn] | None = None
 
     def get_snowpark_column_name_from_spark_column_name(
         self,
@@ -664,20 +760,20 @@ class JoinColumnNameMap(ColumnNameMap):
         *,
         allow_non_exists: bool = False,
         return_first: bool = False,
-        # JoinColumnNameMap will never be called with using columns, so these parameters are not used.
-        is_qualified: bool = False,
-        source_qualifiers: list[str] | None = None,
+        original_snowpark_name: str | None = None,
     ) -> str | None:
         snowpark_column_name_in_left = (
             self.left_column_mapping.get_snowpark_column_name_from_spark_column_name(
                 spark_column_name,
                 allow_non_exists=True,
+                original_snowpark_name=original_snowpark_name,
             )
         )
         snowpark_column_name_in_right = (
             self.right_column_mapping.get_snowpark_column_name_from_spark_column_name(
                 spark_column_name,
                 allow_non_exists=True,
+                original_snowpark_name=original_snowpark_name,
             )
         )
 
@@ -688,14 +784,37 @@ class JoinColumnNameMap(ColumnNameMap):
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Spark column name {spark_column_name} does not exist in either left or right DataFrame"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
 
+        # special case for join conditions, if the column has a match on both sides, and exactly one of those
+        # matches is the original snowpark name, that match should be used
         if (snowpark_column_name_in_right is not None) and (
             snowpark_column_name_in_left is not None
         ):
-            raise AnalysisException(f"Ambiguous column name {spark_column_name}")
+            if (
+                snowpark_column_name_in_left == original_snowpark_name
+                and snowpark_column_name_in_right != original_snowpark_name
+            ):
+                snowpark_column_name_in_right = None
+
+            if (
+                snowpark_column_name_in_right == original_snowpark_name
+                and snowpark_column_name_in_left != original_snowpark_name
+            ):
+                snowpark_column_name_in_left = None
+
+        if (snowpark_column_name_in_right is not None) and (
+            snowpark_column_name_in_left is not None
+        ):
+            exception = AnalysisException(
+                f"Ambiguous column name `{spark_column_name}` in join condition"
+            )
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
 
         snowpark_name = (
             snowpark_column_name_in_right
@@ -703,86 +822,128 @@ class JoinColumnNameMap(ColumnNameMap):
             else snowpark_column_name_in_left
         )
 
-        # this means that the reference is for the column in right dataframe but same snowpark name exist in left dataframe as well
-        # or vice versa, so we need to append _left or _right to the snowpark name
-        if (
-            snowpark_name in self.left_column_mapping.get_snowpark_columns()
-            and snowpark_column_name_in_right is not None
-        ):
-            snowpark_name = quote_name(f"{unquote_if_quoted(snowpark_name)}_right")
-        elif (
-            snowpark_name in self.right_column_mapping.get_snowpark_columns()
-            and snowpark_column_name_in_left is not None
-        ):
-            snowpark_name = quote_name(f"{unquote_if_quoted(snowpark_name)}_left")
-
         return snowpark_name
 
     def get_snowpark_column_names_from_spark_column_names(
-        self, spark_column_names: list[str], return_first: bool = False
+        self,
+        spark_column_names: list[str],
+        return_first: bool = False,
+        original_snowpark_names: list[str] | None = None,
     ) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_column_names_from_snowpark_column_names(
         self,
         snowpark_column_names: list[str],
     ) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_column_name_from_snowpark_column_name(
-        self, snowpark_column_name: str
+        self,
+        snowpark_column_name: str,
+        allow_non_exists: bool = False,
     ) -> str:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_columns(self) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_snowpark_columns(self) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
-    def get_snowpark_columns_after_drop(self, cols_to_drop: list[str]) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+    def get_snowpark_columns_after_drop(
+        self, cols_to_drop: list[str]
+    ) -> list[ColumnNames]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_renamed_nested_column_name(self, name) -> str | None:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def has_spark_column(self, spark_column_name: str) -> bool:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def snowpark_to_spark_map(self) -> dict[str, str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
-    def spark_to_snowpark_for_pattern(self, pattern: str) -> list[tuple[str, str]]:
-        raise NotImplementedError("Method not implemented!")
+    def get_columns_matching_pattern(self, pattern: str) -> list[tuple[str, str]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def with_columns(
         self, new_spark_columns: list[str], new_snowpark_columns: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
-        raise NotImplementedError("Method not implemented!")
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
-    def get_qualifiers(self) -> list[list[str]]:
-        raise NotImplementedError("Method not implemented!")
+    def get_qualifiers(self) -> list[set[ColumnQualifier]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_qualifiers_for_columns_after_drop(
         self, cols_to_drop: list[str]
-    ) -> list[list[str]]:
-        raise NotImplementedError("Method not implemented!")
+    ) -> list[set[ColumnQualifier]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_and_snowpark_columns_with_qualifier_for_qualifier(
-        self, qualifiers_input: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
-        raise NotImplementedError("Method not implemented!")
-
-    def get_qualifier_for_spark_column(self, spark_column_name: str) -> list[str]:
-
-        qualifier_left = self.left_column_mapping.get_qualifier_for_spark_column(
-            spark_column_name
+        self, target_qualifier: list[str]
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
+
+    def get_qualifiers_for_snowpark_column(
+        self, snowpark_name: str
+    ) -> set[ColumnQualifier]:
+        qualifiers_left = self.left_column_mapping.get_qualifiers_for_snowpark_column(
+            snowpark_name
         )
-        qualifier_right = self.right_column_mapping.get_qualifier_for_spark_column(
-            spark_column_name
+        qualifiers_right = self.right_column_mapping.get_qualifiers_for_snowpark_column(
+            snowpark_name
         )
 
-        if (len(qualifier_left) > 0) and (len(qualifier_right) > 0):
-            raise AnalysisException(f"Ambiguous column name {spark_column_name}")
-
-        return qualifier_right if len(qualifier_left) == 0 else qualifier_left
+        if (len(qualifiers_left) > 0) and (len(qualifiers_right) > 0):
+            exception = AnalysisException(f"Ambiguous column name {snowpark_name}")
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
+
+        return qualifiers_right if len(qualifiers_left) == 0 else qualifiers_left
+
+    def get_columns_after_join(
+        self, right: ColumnNameMap, join_columns: list[str], join_type: str
+    ) -> list[ColumnNames]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
+
+    def get_equivalent_snowpark_names_for_snowpark_name(self, snowpark_name: str):
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
+
+    def get_equivalent_snowpark_names(self):
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception