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

snowflake/snowpark_connect/relation/map_join.py

@@ -1,35 +1,50 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-from collections import Counter
+import dataclasses
+from collections.abc import Callable
+from copy import copy
+from enum import Enum
 from functools import reduce
+from typing import Optional
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
-from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.errors import AnalysisException
+from pyspark.errors.exceptions.connect import IllegalArgumentException
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
-from snowflake.snowpark._internal.analyzer.analyzer_utils import (
-    quote_name_without_upper_casing,
-    unquote_if_quoted,
+from snowflake.snowpark import Column, DataFrame
+from snowflake.snowpark.types import StructField, StructType
+from snowflake.snowpark_connect.column_name_handler import (
+    ColumnNames,
+    ColumnQualifier,
+    JoinColumnNameMap,
+    make_unique_snowpark_name,
 )
-from snowflake.snowpark_connect.column_name_handler import JoinColumnNameMap
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.constants import COLUMN_METADATA_COLLISION_KEY
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
-from snowflake.snowpark_connect.error.error_utils import SparkException
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import (
+    SparkException,
+    attach_custom_error_code,
+)
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
 )
 from snowflake.snowpark_connect.expression.typer import JoinExpressionTyper
-from snowflake.snowpark_connect.hidden_column import HiddenColumn
 from snowflake.snowpark_connect.relation.map_relation import (
     NATURAL_JOIN_TYPE_BASE,
     map_relation,
 )
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    without_internal_columns,
+)
 from snowflake.snowpark_connect.utils.context import (
     push_evaluating_join_condition,
     push_sql_scope,
+    set_plan_id_map,
     set_sql_plan_name,
 )
 from snowflake.snowpark_connect.utils.telemetry import (
@@ -38,447 +53,583 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 USING_COLUMN_NOT_FOUND_ERROR = "[UNRESOLVED_USING_COLUMN_FOR_JOIN] USING column `{0}` not found on the {1} side of the join. The {1}-side columns: {2}"
 
-DUPLICATED_JOIN_COL_LSUFFIX = "_left"
-DUPLICATED_JOIN_COL_RSUFFIX = "_right"
+
+class ConditionType(Enum):
+    USING_COLUMNS = 1
+    JOIN_CONDITION = 2
+    NO_CONDITION = 3
+
+
+@dataclasses.dataclass
+class JoinInfo:
+    join_type: str
+    condition_type: ConditionType
+    join_columns: Optional[list[str]]
+    just_left_columns: bool
+    is_join_with: bool
+    is_left_struct: bool
+    is_right_struct: bool
+
+    def is_using_columns(self):
+        return self.condition_type == ConditionType.USING_COLUMNS
 
 
 def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
     left_container: DataFrameContainer = map_relation(rel.join.left)
     right_container: DataFrameContainer = map_relation(rel.join.right)
 
-    left_input: snowpark.DataFrame = left_container.dataframe
-    right_input: snowpark.DataFrame = right_container.dataframe
-    is_natural_join = rel.join.join_type >= NATURAL_JOIN_TYPE_BASE
-    using_columns = rel.join.using_columns
-    if is_natural_join:
-        rel.join.join_type -= NATURAL_JOIN_TYPE_BASE
-        left_spark_columns = left_container.column_map.get_spark_columns()
-        right_spark_columns = right_container.column_map.get_spark_columns()
-        common_spark_columns = [
-            x for x in left_spark_columns if x in right_spark_columns
-        ]
-        using_columns = common_spark_columns
+    # Remove any metadata columns(like metada$filename) present in the dataframes.
+    # We cannot support inputfilename for multisources as each dataframe has it's own source.
+    left_container = without_internal_columns(left_container)
+    right_container = without_internal_columns(right_container)
 
-    match rel.join.join_type:
-        case relation_proto.Join.JOIN_TYPE_UNSPECIFIED:
-            # TODO: Understand what UNSPECIFIED Join type is
-            raise SnowparkConnectNotImplementedError("Unspecified Join Type")
-        case relation_proto.Join.JOIN_TYPE_INNER:
-            join_type = "inner"
-        case relation_proto.Join.JOIN_TYPE_FULL_OUTER:
-            join_type = "full_outer"
-        case relation_proto.Join.JOIN_TYPE_LEFT_OUTER:
-            join_type = "left"
-        case relation_proto.Join.JOIN_TYPE_RIGHT_OUTER:
-            join_type = "right"
-        case relation_proto.Join.JOIN_TYPE_LEFT_ANTI:
-            join_type = "leftanti"
-        case relation_proto.Join.JOIN_TYPE_LEFT_SEMI:
-            join_type = "leftsemi"
-        case relation_proto.Join.JOIN_TYPE_CROSS:
-            join_type = "cross"
-        case other:
-            raise SnowparkConnectNotImplementedError(f"Other Join Type: {other}")
-
-    # This handles case sensitivity for using_columns
-    case_corrected_right_columns: list[str] = []
-    hidden_columns = set()
-    # Propagate the hidden columns from left/right inputs to the result in case of chained joins
-    if left_container.column_map.hidden_columns:
-        hidden_columns.update(left_container.column_map.hidden_columns)
-
-    if right_container.column_map.hidden_columns:
-        hidden_columns.update(right_container.column_map.hidden_columns)
-
-    if rel.join.HasField("join_condition"):
-        assert not using_columns
-
-        left_columns = list(left_container.column_map.spark_to_col.keys())
-        right_columns = list(right_container.column_map.spark_to_col.keys())
-
-        # All PySpark join types are in the format of JOIN_TYPE_XXX.
-        # We remove the first 10 characters (JOIN_TYPE_) and replace all underscores with spaces to match the exception.
-        pyspark_join_type = relation_proto.Join.JoinType.Name(rel.join.join_type)[
-            10:
-        ].replace("_", " ")
-        with push_sql_scope(), push_evaluating_join_condition(
-            pyspark_join_type, left_columns, right_columns
-        ):
-            if left_container.alias is not None:
-                set_sql_plan_name(left_container.alias, rel.join.left.common.plan_id)
-            if right_container.alias is not None:
-                set_sql_plan_name(right_container.alias, rel.join.right.common.plan_id)
-            _, join_expression = map_single_column_expression(
-                rel.join.join_condition,
-                column_mapping=JoinColumnNameMap(
-                    left_container.column_map,
-                    right_container.column_map,
-                ),
-                typer=JoinExpressionTyper(left_input, right_input),
+    left_plan = rel.join.left.common.plan_id
+    right_plan = rel.join.right.common.plan_id
+
+    # if there are any conflicting snowpark columns, this is the time to rename them
+    disambiguated_right_container = _disambiguate_snowpark_columns(
+        left_container, right_container, right_plan if left_plan != right_plan else None
+    )
+
+    join_info = _get_join_info(rel, left_container, disambiguated_right_container)
+
+    match join_info.condition_type:
+        case ConditionType.JOIN_CONDITION:
+            result_container = _join_using_condition(
+                left_container,
+                disambiguated_right_container,
+                join_info,
+                rel,
+                right_container if left_plan == right_plan else None,
             )
-        result: snowpark.DataFrame = left_input.join(
-            right=right_input,
-            on=join_expression.col,
-            how=join_type,
-            lsuffix=DUPLICATED_JOIN_COL_LSUFFIX,
-            rsuffix=DUPLICATED_JOIN_COL_RSUFFIX,
-        )
-    elif using_columns:
-        if any(
-            left_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                c, allow_non_exists=True, return_first=True
+        case ConditionType.USING_COLUMNS:
+            result_container = _join_using_columns(
+                left_container,
+                disambiguated_right_container,
+                join_info,
             )
-            is None
-            for c in using_columns
-        ):
-            import pyspark
-
-            raise pyspark.errors.AnalysisException(
-                USING_COLUMN_NOT_FOUND_ERROR.format(
-                    next(
-                        c
-                        for c in using_columns
-                        if left_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                            c, allow_non_exists=True, return_first=True
-                        )
-                        is None
-                    ),
-                    "left",
-                    left_container.column_map.get_spark_columns(),
-                )
+        case _:
+            result_container = _join_unconditionally(
+                left_container, disambiguated_right_container, join_info
             )
-        if any(
-            right_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                c, allow_non_exists=True, return_first=True
+
+    return result_container
+
+
+def _join_unconditionally(
+    left_container: DataFrameContainer,
+    right_container: DataFrameContainer,
+    info: JoinInfo,
+) -> DataFrameContainer:
+    if info.join_type != "cross" and not global_config.spark_sql_crossJoin_enabled:
+        exception = SparkException.implicit_cartesian_product("inner")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
+
+    left_input = left_container.dataframe
+    right_input = right_container.dataframe
+    join_type = info.join_type
+
+    # For outer joins without a condition, we need to use a TRUE condition
+    # to match Spark's behavior.
+    result: snowpark.DataFrame = left_input.join(
+        right=right_input,
+        on=snowpark_fn.lit(True)
+        if join_type in ["left", "right", "full_outer"]
+        else None,
+        how=join_type,
+    )
+
+    columns = left_container.column_map.columns + right_container.column_map.columns
+    column_metadata = _combine_metadata(left_container, right_container)
+
+    if info.just_left_columns:
+        columns = left_container.column_map.columns
+        column_metadata = left_container.column_map.column_metadata
+        result = result.select(*left_container.column_map.get_snowpark_columns())
+
+    snowpark_columns = [c.snowpark_name for c in columns]
+
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=[c.spark_name for c in columns],
+        snowpark_column_names=snowpark_columns,
+        column_metadata=column_metadata,
+        column_qualifiers=[c.qualifiers for c in columns],
+        cached_schema_getter=_build_joined_schema(
+            snowpark_columns, left_input, right_input
+        ),
+        equivalent_snowpark_names=[c.equivalent_snowpark_names for c in columns],
+    )
+
+
+def _join_using_columns(
+    left_container: DataFrameContainer,
+    right_container: DataFrameContainer,
+    info: JoinInfo,
+) -> DataFrameContainer:
+    join_columns = info.join_columns
+
+    def _validate_using_column(
+        column: str, container: DataFrameContainer, side: str
+    ) -> None:
+        if (
+            container.column_map.get_snowpark_column_name_from_spark_column_name(
+                column, allow_non_exists=True, return_first=True
             )
             is None
-            for c in using_columns
         ):
-            import pyspark
-
-            raise pyspark.errors.AnalysisException(
+            exception = AnalysisException(
                 USING_COLUMN_NOT_FOUND_ERROR.format(
-                    next(
-                        c
-                        for c in using_columns
-                        if right_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                            c, allow_non_exists=True, return_first=True
-                        )
-                        is None
-                    ),
-                    "right",
-                    right_container.column_map.get_spark_columns(),
+                    column, side, container.column_map.get_spark_columns()
                 )
             )
-
-        using_columns_snowpark_names = (
-            left_container.column_map.get_snowpark_column_names_from_spark_column_names(
-                list(using_columns), return_first=True
-            )
+            attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+            raise exception
+
+    for col in join_columns:
+        _validate_using_column(col, left_container, "left")
+        _validate_using_column(col, right_container, "right")
+
+    left_input = left_container.dataframe
+    right_input = right_container.dataframe
+
+    # The inputs will have different snowpark names for the same spark name,
+    # so we convert ["a", "b"] into (left["a"] == right["a"] & left["b"] == right["b"]),
+    # then drop right["a"] and right["b"].
+    snowpark_using_columns = [
+        (
+            snowpark_fn.col(
+                left_container.column_map.get_snowpark_column_name_from_spark_column_name(
+                    spark_name, return_first=True
+                )
+            ),
+            snowpark_fn.col(
+                right_container.column_map.get_snowpark_column_name_from_spark_column_name(
+                    spark_name, return_first=True
+                )
+            ),
         )
+        for spark_name in join_columns
+    ]
+
+    # this is a condition join, so it will contain left + right columns
+    # we need to postprocess this later to have a correct projection
+    joined_df = left_input.join(
+        right=right_input,
+        on=reduce(
+            snowpark.Column.__and__,
+            (left == right for left, right in snowpark_using_columns),
+        ),
+        how=info.join_type,
+    )
 
-        using_columns_snowpark_types = [
-            left_container.dataframe.schema.fields[idx].datatype
-            for idx, col in enumerate(left_container.column_map.get_snowpark_columns())
-            if col in using_columns_snowpark_names
-        ]
+    # figure out default column ordering after the join
+    columns = left_container.column_map.get_columns_after_join(
+        right_container.column_map, join_columns, info.join_type
+    )
+
+    if info.join_type in ["full_outer", "left", "right"]:
+        all_columns_for_select = []
+        all_column_names = []
 
-        # Round trip the using columns through the column map to get the correct names
-        # in order to support case sensitivity.
-        # TODO: case_corrected_left_columns / case_corrected_right_columns may no longer be required as Snowpark dataframe preserves the column casing now.
-        case_corrected_left_columns = (
-            left_container.column_map.get_spark_column_names_from_snowpark_column_names(
-                using_columns_snowpark_names
+        for column_info in columns[: len(join_columns)]:
+            spark_name = column_info.spark_name
+            left_sp_name = left_container.column_map.get_snowpark_column_name_from_spark_column_name(
+                spark_name, return_first=True
             )
-        )
-        case_corrected_right_columns = right_container.column_map.get_spark_column_names_from_snowpark_column_names(
-            right_container.column_map.get_snowpark_column_names_from_spark_column_names(
-                list(using_columns), return_first=True
+            right_sp_name = right_container.column_map.get_snowpark_column_name_from_spark_column_name(
+                spark_name, return_first=True
             )
-        )
-        using_columns = zip(case_corrected_left_columns, case_corrected_right_columns)
-        # We cannot assume that Snowpark will have the same names for left and right columns,
-        # so we convert ["a", "b"] into (left["a"] == right["a"] & left["b"] == right["b"]),
-        # then drop right["a"] and right["b"].
-        snowpark_using_columns = [
-            (
-                left_input[
-                    left_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                        lft, return_first=True
-                    )
-                ],
-                right_input[
-                    right_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                        r, return_first=True
+
+            if info.join_type == "full_outer":
+                new_sp_name = make_unique_snowpark_name(spark_name)
+                all_columns_for_select.append(
+                    snowpark_fn.coalesce(
+                        snowpark_fn.col(left_sp_name), snowpark_fn.col(right_sp_name)
+                    ).alias(new_sp_name)
+                )
+                all_column_names.append(
+                    ColumnNames(
+                        spark_name,
+                        new_sp_name,
+                        set(),
+                        equivalent_snowpark_names=set(),
+                        is_hidden=False,
                     )
-                ],
-            )
-            for lft, r in using_columns
-        ]
-        joined_df = left_input.join(
-            right=right_input,
-            on=reduce(
-                snowpark.Column.__and__,
-                (left == right for left, right in snowpark_using_columns),
-            ),
-            how=join_type,
-            rsuffix=DUPLICATED_JOIN_COL_RSUFFIX,
-        )
-        # If we disambiguated the snowpark_using_columns during the join, we need to update 'snowpark_using_columns' to
-        # use the disambiguated names.
-        disambiguated_snowpark_using_columns = []
+                )
 
-        # Ignore disambiguation for LEFT SEMI JOIN and LEFT ANTI JOIN because they drop the right columns, so it'll never disambiguate.
-        if join_type in ["leftsemi", "leftanti"]:
-            disambiguated_snowpark_using_columns = snowpark_using_columns
-        else:
-            normalized_joined_columns = [
-                unquote_if_quoted(col) for col in joined_df.columns
-            ]
-            # snowpark_using_columns is a list of tuples of snowpark columns, joined_df.columns is a list of strings of column names
-            for (left, right) in snowpark_using_columns:
-                normalized_left_name = unquote_if_quoted(left.getName())
-                normalized_right_name = unquote_if_quoted(right.getName())
-
-                # are both left and right in joined_df? if not, it's been disambiguated
-                if (
-                    normalized_left_name in normalized_joined_columns
-                    and normalized_right_name in normalized_joined_columns
-                ):
-                    # we want to just add this
-                    disambiguated_snowpark_using_columns.append((left, right))
-                else:
-                    # we need to figure out the disambiguated names and add those - it only disambiguates if left == right
-                    disambiguated_left: snowpark.Column | None = None
-                    disambiguated_right: snowpark.Column | None = None
-
-                    for col in normalized_joined_columns:
-                        quoted_col = f'"{col}"'
-                        # get the column name and cross check it to see if it ends with the og name
-                        if col.endswith(normalized_left_name) and col.startswith("l_"):
-                            disambiguated_left = joined_df[quoted_col]
-                        elif col.endswith(normalized_right_name) and col.startswith(
-                            "r_"
-                        ):
-                            disambiguated_right = joined_df[quoted_col]
-
-                        # If we have both disambiguated columns, we can break out of the loop to save processing time
-                        if (
-                            disambiguated_left is not None
-                            and disambiguated_right is not None
-                        ):
-                            break
-                    if disambiguated_left is None or disambiguated_right is None:
-                        raise AnalysisException(
-                            f"Disambiguated columns not found for {normalized_left_name} and {normalized_right_name}."
+                for sp_name, container in [
+                    (left_sp_name, left_container),
+                    (right_sp_name, right_container),
+                ]:
+                    all_columns_for_select.append(snowpark_fn.col(sp_name))
+                    all_column_names.append(
+                        ColumnNames(
+                            spark_name,
+                            sp_name,
+                            container.column_map.get_qualifiers_for_snowpark_column(
+                                sp_name
+                            ),
+                            equivalent_snowpark_names=container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                                sp_name
+                            ),
+                            is_hidden=True,
+                        )
+                    )
+            else:
+                for sp_name, container, side in [
+                    (left_sp_name, left_container, "left"),
+                    (right_sp_name, right_container, "right"),
+                ]:
+                    all_columns_for_select.append(snowpark_fn.col(sp_name))
+                    qualifiers = (
+                        container.column_map.get_qualifiers_for_snowpark_column(sp_name)
+                    )
+                    equivalent_snowpark_names = set()
+                    equivalent_snowpark_names.update(
+                        container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                            sp_name
+                        )
+                    )
+                    is_visible = info.join_type == side
+                    if is_visible:
+                        qualifiers = qualifiers | {ColumnQualifier(())}
+                    all_column_names.append(
+                        ColumnNames(
+                            spark_name,
+                            sp_name,
+                            qualifiers,
+                            equivalent_snowpark_names=equivalent_snowpark_names,
+                            is_hidden=not is_visible,
                         )
-                    disambiguated_snowpark_using_columns.append(
-                        (disambiguated_left, disambiguated_right)
                     )
 
-        # For outer joins, we need to preserve join keys from both sides using COALESCE
-        """
-        CHANGES:
-            - IF CASE
-                - Need to drop the using columns
-                - Need to create the hidden_columns DF with the using columns from right and left
-            - ELSE CASE
-                - Need to drop the right side using columns
-                - Need to create the hidden_columns DF with the using columns from right
-        """
-        if join_type == "full_outer":
-            coalesced_columns = []
-            for i, (left_col, _right_col) in enumerate(snowpark_using_columns):
-                # Use the original user-specified column name to preserve case sensitivity
-                # Use the disambiguated columns for coalescing
-                disambiguated_left_col = disambiguated_snowpark_using_columns[i][0]
-                disambiguated_right_col = disambiguated_snowpark_using_columns[i][1]
-
-                coalesced_col = snowpark_fn.coalesce(
-                    disambiguated_left_col, disambiguated_right_col
-                ).alias(left_col.get_name())
-                coalesced_columns.append(coalesced_col)
-
-                # Create HiddenColumn objects for each hidden column
-                hidden_left = HiddenColumn(
-                    hidden_snowpark_name=disambiguated_left_col.getName(),
-                    spark_name=case_corrected_left_columns[i],
-                    visible_snowpark_name=left_col.get_name(),
-                    qualifiers=left_container.column_map.get_qualifier_for_spark_column(
-                        case_corrected_left_columns[i]
-                    ),
-                    original_position=left_container.column_map.get_spark_columns().index(
-                        case_corrected_left_columns[i]
-                    ),
-                )
-
-                hidden_right = HiddenColumn(
-                    hidden_snowpark_name=disambiguated_right_col.getName(),
-                    spark_name=case_corrected_right_columns[i],
-                    visible_snowpark_name=left_col.get_name(),
-                    qualifiers=right_container.column_map.get_qualifier_for_spark_column(
-                        case_corrected_right_columns[i]
-                    ),
-                    original_position=right_container.column_map.get_spark_columns().index(
-                        case_corrected_right_columns[i]
-                    ),
-                )
-                hidden_columns.update(
-                    [
-                        hidden_left,
-                        hidden_right,
-                    ]
-                )
+        for c in columns[len(join_columns) :]:
+            all_columns_for_select.append(snowpark_fn.col(c.snowpark_name))
+            all_column_names.append(c)
 
-            # All non-hidden columns (not including the coalesced columns)
-            other_columns = [
-                snowpark_fn.col(col_name)
-                for col_name in joined_df.columns
-                if col_name not in [col.hidden_snowpark_name for col in hidden_columns]
-            ]
-            result = joined_df.select(coalesced_columns + other_columns)
+        result = joined_df.select(all_columns_for_select)
+        snowpark_names_for_schema = [c.snowpark_name for c in columns]
 
-        else:
-            result = joined_df.drop(*(right for _, right in snowpark_using_columns))
-            # We never run into the disambiguation case unless it's a full outer join.
-            for i, (left_col, right_col) in enumerate(
-                disambiguated_snowpark_using_columns
-            ):
-                # Only right side columns are hidden
-                hidden_col = HiddenColumn(
-                    hidden_snowpark_name=right_col.getName(),
-                    spark_name=case_corrected_right_columns[i],
-                    visible_snowpark_name=left_col.getName(),
-                    qualifiers=right_container.column_map.get_qualifier_for_spark_column(
-                        case_corrected_right_columns[i]
-                    ),
-                    original_position=right_container.column_map.get_spark_columns().index(
-                        case_corrected_right_columns[i]
-                    ),
-                )
-                hidden_columns.add(hidden_col)
-    else:
-        if join_type != "cross" and not global_config.spark_sql_crossJoin_enabled:
-            raise SparkException.implicit_cartesian_product("inner")
-        result: snowpark.DataFrame = left_input.join(
-            right=right_input,
-            how=join_type,
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result,
+            spark_column_names=[c.spark_name for c in all_column_names],
+            snowpark_column_names=[c.snowpark_name for c in all_column_names],
+            column_metadata=_combine_metadata(left_container, right_container),
+            column_qualifiers=[c.qualifiers for c in all_column_names],
+            column_is_hidden=[c.is_hidden for c in all_column_names],
+            cached_schema_getter=_build_joined_schema(
+                snowpark_names_for_schema,
+                left_input,
+                right_input,
+                all_column_names,
+            ),
+            equivalent_snowpark_names=[
+                c.equivalent_snowpark_names for c in all_column_names
+            ],
         )
 
-    if join_type in ["leftanti", "leftsemi"]:
-        # Join types that only return columns from the left side:
-        # - LEFT SEMI JOIN: Returns left rows that have matches in right table (no right columns)
-        # - LEFT ANTI JOIN: Returns left rows that have NO matches in right table (no right columns)
-        # Both preserve only the columns from the left DataFrame without adding any columns from the right.
-        spark_cols_after_join = left_container.column_map.get_spark_columns()
-        snowpark_cols_after_join = left_container.column_map.get_snowpark_columns()
-        snowpark_col_types = [
-            f.datatype for f in left_container.dataframe.schema.fields
-        ]
-        qualifiers = left_container.column_map.get_qualifiers()
-    elif join_type == "full_outer" and using_columns:
-        # We want the coalesced columns to be first, followed by all the left and right columns (excluding using columns)
-        spark_cols_after_join: list[str] = []
-        snowpark_cols_after_join: list[str] = []
-        snowpark_col_types: list[str] = []
-
-        left_container_snowpark_columns = (
-            left_container.column_map.get_snowpark_columns()
+    if info.just_left_columns:
+        # we just need the left columns
+        columns = columns[: len(left_container.column_map.columns)]
+        snowpark_columns = [c.snowpark_name for c in columns]
+        result = joined_df.select(*snowpark_columns)
+
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result,
+            spark_column_names=[c.spark_name for c in columns],
+            snowpark_column_names=snowpark_columns,
+            column_metadata=left_container.column_map.column_metadata,
+            column_qualifiers=[c.qualifiers for c in columns],
+            cached_schema_getter=_build_joined_schema(
+                snowpark_columns, left_input, right_input
+            ),
+            equivalent_snowpark_names=[c.equivalent_snowpark_names for c in columns],
         )
-        right_container_snowpark_columns = (
-            right_container.column_map.get_snowpark_columns()
+
+    snowpark_columns = [c.snowpark_name for c in columns]
+    result = joined_df.select(*snowpark_columns)
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=[c.spark_name for c in columns],
+        snowpark_column_names=snowpark_columns,
+        column_metadata=_combine_metadata(left_container, right_container),
+        column_qualifiers=[c.qualifiers for c in columns],
+        cached_schema_getter=_build_joined_schema(
+            snowpark_columns, left_input, right_input
+        ),
+        equivalent_snowpark_names=[c.equivalent_snowpark_names for c in columns],
+    )
+
+
+def _join_using_condition(
+    left_container: DataFrameContainer,
+    right_container: DataFrameContainer,
+    info: JoinInfo,
+    rel: relation_proto.Relation,
+    original_right_container: Optional[DataFrameContainer],
+) -> DataFrameContainer:
+    left_columns = left_container.column_map.get_spark_columns()
+    right_columns = right_container.column_map.get_spark_columns()
+
+    left_input = left_container.dataframe
+    right_input = right_container.dataframe
+
+    # All PySpark join types are in the format of JOIN_TYPE_XXX.
+    # We remove the first 10 characters (JOIN_TYPE_) and replace all underscores with spaces to match the exception.
+    pyspark_join_type = relation_proto.Join.JoinType.Name(rel.join.join_type)[
+        10:
+    ].replace("_", " ")
+    with push_sql_scope(), push_evaluating_join_condition(
+        pyspark_join_type, left_columns, right_columns
+    ):
+        if left_container.alias is not None:
+            set_sql_plan_name(left_container.alias, rel.join.left.common.plan_id)
+        if right_container.alias is not None:
+            set_sql_plan_name(right_container.alias, rel.join.right.common.plan_id)
+        # resolve join condition expression
+        _, join_expression = map_single_column_expression(
+            rel.join.join_condition,
+            column_mapping=JoinColumnNameMap(
+                left_container.column_map,
+                # using the original (not disambiguated) right container is intended to break
+                # self join cases like a.join(a, a.id == a.id), since SAS can't handle them correctly
+                # and they fail in Spark Connect
+                (
+                    original_right_container
+                    if original_right_container
+                    else right_container
+                ).column_map,
+            ),
+            typer=JoinExpressionTyper(left_input, right_input),
         )
 
-        qualifiers = []
-        for i in range(len(case_corrected_left_columns)):
-            spark_cols_after_join.append(case_corrected_left_columns[i])
-            snowpark_cols_after_join.append(using_columns_snowpark_names[i])
-            snowpark_col_types.append(using_columns_snowpark_types[i])
-            qualifiers.append([])
-
-        # Handle adding left and right columns, excluding the using columns
-        for i, spark_col in enumerate(left_container.column_map.get_spark_columns()):
-            if (
-                spark_col not in case_corrected_left_columns
-                or spark_col in left_container.column_map.get_spark_columns()[:i]
-            ):
-                spark_cols_after_join.append(spark_col)
-                snowpark_cols_after_join.append(left_container_snowpark_columns[i])
-                qualifiers.append(
-                    left_container.column_map.get_qualifier_for_spark_column(spark_col)
-                )
+    result: snowpark.DataFrame = left_input.join(
+        right=right_input,
+        on=join_expression.col,
+        how=info.join_type,
+    )
 
-                snowpark_col_types.append(
-                    left_container.dataframe.schema.fields[i].datatype
-                )
+    # early return for joinWith
+    if info.is_join_with:
+        return _join_with(left_container, right_container, result, info)
 
-        for i, spark_col in enumerate(right_container.column_map.get_spark_columns()):
-            if (
-                spark_col not in case_corrected_right_columns
-                or spark_col in right_container.column_map.get_spark_columns()[:i]
-            ):
-                spark_cols_after_join.append(spark_col)
-                snowpark_cols_after_join.append(right_container_snowpark_columns[i])
-                qualifiers.append(
-                    right_container.column_map.get_qualifier_for_spark_column(spark_col)
-                )
+    # column order is already correct, so we just take the left + right side list
+    columns = left_container.column_map.columns + right_container.column_map.columns
+    column_metadata = _combine_metadata(left_container, right_container)
 
-                snowpark_col_types.append(
-                    right_container.dataframe.schema.fields[i].datatype
-                )
+    if info.just_left_columns:
+        # we just need left-side columns
+        columns = left_container.column_map.columns
+        result = result.select(*[c.snowpark_name for c in columns])
+        column_metadata = left_container.column_map.column_metadata
 
-    else:
-        spark_cols_after_join = left_container.column_map.get_spark_columns()
-        snowpark_cols_after_join = left_container.column_map.get_snowpark_columns()
-        snowpark_col_types = [
-            f.datatype for f in left_container.dataframe.schema.fields
-        ]
+    snowpark_columns = [c.snowpark_name for c in columns]
 
-        qualifiers = left_container.column_map.get_qualifiers()
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=[c.spark_name for c in columns],
+        snowpark_column_names=snowpark_columns,
+        column_metadata=column_metadata,
+        column_qualifiers=[c.qualifiers for c in columns],
+        cached_schema_getter=_build_joined_schema(
+            snowpark_columns, left_input, right_input
+        ),
+        equivalent_snowpark_names=[c.equivalent_snowpark_names for c in columns],
+    )
 
-        right_df_snowpark_columns = right_container.column_map.get_snowpark_columns()
 
-        for i, spark_col in enumerate(right_container.column_map.get_spark_columns()):
-            if (
-                spark_col not in case_corrected_right_columns
-                or spark_col in right_container.column_map.get_spark_columns()[:i]
-            ):
-                spark_cols_after_join.append(spark_col)
-                snowpark_cols_after_join.append(right_df_snowpark_columns[i])
-                snowpark_col_types.append(
-                    right_container.dataframe.schema.fields[i].datatype
-                )
+def _join_with(
+    left_container: DataFrameContainer,
+    right_container: DataFrameContainer,
+    joined_df: DataFrame,
+    info: JoinInfo,
+) -> DataFrameContainer:
+    # joinWith always returns 2 columns
+    left_column = "_1"
+    right_column = "_2"
+    left_snowpark_name: str = make_unique_snowpark_name(left_column)
+    right_snowpark_name: str = make_unique_snowpark_name(right_column)
 
-                qualifiers.append(
-                    right_container.column_map.get_qualifier_for_spark_column(spark_col)
-                )
+    left_nullable, right_nullable = _join_with_nullability(info.join_type)
 
-    snowpark_cols_after_join_deduplicated = []
-    snowpark_cols_after_join_counter = Counter(snowpark_cols_after_join)
-    seen_duplicated_columns = set()
+    left_col, left_col_type = _construct_join_with_column(
+        left_container, left_snowpark_name, info.is_left_struct
+    )
+    right_col, right_col_type = _construct_join_with_column(
+        right_container, right_snowpark_name, info.is_right_struct
+    )
 
-    for col in snowpark_cols_after_join:
-        if snowpark_cols_after_join_counter[col] == 2:
-            # This means that the same column exists twice in the joined df, likely due to a self-join and
-            # we need to lsuffix and rsuffix to the names of both columns, similar to what Snowpark did under the hood.
+    result = joined_df.select(left_col, right_col)
 
-            suffix = (
-                DUPLICATED_JOIN_COL_RSUFFIX
-                if col in seen_duplicated_columns
-                else DUPLICATED_JOIN_COL_LSUFFIX
-            )
-            unquoted_col = unquote_if_quoted(col)
-            quoted = quote_name_without_upper_casing(unquoted_col + suffix)
-            snowpark_cols_after_join_deduplicated.append(quoted)
+    def _schema_getter() -> StructType:
+        return StructType(
+            [
+                StructField(left_snowpark_name, left_col_type, left_nullable),
+                StructField(right_snowpark_name, right_col_type, right_nullable),
+            ]
+        )
+
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=[left_column, right_column],
+        snowpark_column_names=[left_snowpark_name, right_snowpark_name],
+        cached_schema_getter=_schema_getter,
+        column_metadata={},  # no top-level metadata for struct columns
+        # no qualifiers or equivalent snowpark names
+    )
+
+
+def _get_join_info(
+    rel: relation_proto.Relation, left: DataFrameContainer, right: DataFrameContainer
+) -> JoinInfo:
+    """
+    Gathers basic information about the join, and performs basic assertions
+    """
+
+    is_natural_join = rel.join.join_type >= NATURAL_JOIN_TYPE_BASE
+    join_columns = rel.join.using_columns
+    if is_natural_join:
+        rel.join.join_type -= NATURAL_JOIN_TYPE_BASE
+        left_spark_columns = left.column_map.get_spark_columns()
+        right_spark_columns = right.column_map.get_spark_columns()
+        common_spark_columns = [
+            x for x in left_spark_columns if x in right_spark_columns
+        ]
+        join_columns = common_spark_columns
 
-            seen_duplicated_columns.add(col)
+    match rel.join.join_type:
+        case relation_proto.Join.JOIN_TYPE_UNSPECIFIED:
+            # TODO: Understand what UNSPECIFIED Join type is
+            exception = SnowparkConnectNotImplementedError("Unspecified Join Type")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
+        case relation_proto.Join.JOIN_TYPE_INNER:
+            join_type = "inner"
+        case relation_proto.Join.JOIN_TYPE_FULL_OUTER:
+            join_type = "full_outer"
+        case relation_proto.Join.JOIN_TYPE_LEFT_OUTER:
+            join_type = "left"
+        case relation_proto.Join.JOIN_TYPE_RIGHT_OUTER:
+            join_type = "right"
+        case relation_proto.Join.JOIN_TYPE_LEFT_ANTI:
+            join_type = "leftanti"
+        case relation_proto.Join.JOIN_TYPE_LEFT_SEMI:
+            join_type = "leftsemi"
+        case relation_proto.Join.JOIN_TYPE_CROSS:
+            join_type = "cross"
+        case other:
+            exception = SnowparkConnectNotImplementedError(f"Other Join Type: {other}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
+
+    has_join_condition = rel.join.HasField("join_condition")
+    is_using_columns = bool(join_columns)
+
+    if join_type == "cross" and has_join_condition:
+        # if the user provided any condition, it's no longer a cross join
+        join_type = "inner"
+
+    if has_join_condition:
+        assert not is_using_columns
+
+    condition_type = ConditionType.NO_CONDITION
+    if has_join_condition:
+        condition_type = ConditionType.JOIN_CONDITION
+    elif is_using_columns:
+        condition_type = ConditionType.USING_COLUMNS
+
+    # Join types that only return columns from the left side:
+    # - LEFT SEMI JOIN: Returns left rows that have matches in right table (no right columns)
+    # - LEFT ANTI JOIN: Returns left rows that have NO matches in right table (no right columns)
+    # Both preserve only the columns from the left DataFrame without adding any columns from the right.
+    just_left_columns = join_type in ["leftanti", "leftsemi"]
+
+    # joinWith
+    is_join_with = rel.join.HasField("join_data_type")
+    is_left_struct = False
+    is_right_struct = False
+    if is_join_with:
+        is_left_struct = rel.join.join_data_type.is_left_struct
+        is_right_struct = rel.join.join_data_type.is_right_struct
+
+    return JoinInfo(
+        join_type,
+        condition_type,
+        join_columns,
+        just_left_columns,
+        is_join_with,
+        is_left_struct,
+        is_right_struct,
+    )
+
+
+def _disambiguate_snowpark_columns(
+    left: DataFrameContainer, right: DataFrameContainer, right_plan: int
+) -> DataFrameContainer:
+    conflicting_snowpark_columns = left.column_map.get_conflicting_snowpark_columns(
+        right.column_map
+    )
+
+    if not conflicting_snowpark_columns:
+        return right
+
+    # rename and create new right container
+    column_map = right.column_map
+    disambiguated_columns: list[Column] = []
+    disambiguated_snowpark_names: list[str] = []
+    # retain old snowpark names in column map
+    equivalent_snowpark_names: list[set[str]] = []
+    for c in column_map.columns:
+        col_equivalent_snowpark_names = copy(c.equivalent_snowpark_names)
+        if c.snowpark_name in conflicting_snowpark_columns:
+            # alias snowpark column with a new unique name
+            new_name = make_unique_snowpark_name(c.spark_name)
+            disambiguated_snowpark_names.append(new_name)
+            disambiguated_columns.append(
+                snowpark_fn.col(c.snowpark_name).alias(new_name)
+            )
         else:
-            snowpark_cols_after_join_deduplicated.append(col)
+            disambiguated_snowpark_names.append(c.snowpark_name)
+            disambiguated_columns.append(snowpark_fn.col(c.snowpark_name))
+
+        equivalent_snowpark_names.append(col_equivalent_snowpark_names)
+
+    disambiguated_df = right.dataframe.select(*disambiguated_columns)
+
+    def _schema_getter() -> StructType:
+        fields = right.dataframe.schema.fields
+        return StructType(
+            [
+                StructField(name, fields[i].datatype, fields[i].nullable)
+                for i, name in enumerate(disambiguated_snowpark_names)
+            ]
+        )
+
+    disambiguated_right = DataFrameContainer.create_with_column_mapping(
+        dataframe=disambiguated_df,
+        spark_column_names=column_map.get_spark_columns(),
+        snowpark_column_names=disambiguated_snowpark_names,
+        column_metadata=column_map.column_metadata,
+        column_qualifiers=column_map.get_qualifiers(),
+        table_name=right.table_name,
+        cached_schema_getter=_schema_getter,
+        equivalent_snowpark_names=equivalent_snowpark_names,
+    )
 
-    column_metadata = {}
-    if left_container.column_map.column_metadata:
-        column_metadata.update(left_container.column_map.column_metadata)
+    # since we just renamed some snowpark columns, we need to update the dataframe container for the given plan_id
+    # TODO: is there a better way to do this?
+    if right_plan:
+        set_plan_id_map(right_plan, disambiguated_right)
 
+    return disambiguated_right
+
+
+def _combine_metadata(
+    left_container: DataFrameContainer, right_container: DataFrameContainer
+) -> dict:
+    column_metadata = dict(left_container.column_map.column_metadata or {})
     if right_container.column_map.column_metadata:
         for key, value in right_container.column_map.column_metadata.items():
             if key not in column_metadata:
@@ -490,7 +641,9 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
                     snowpark_name = right_container.column_map.get_snowpark_column_name_from_spark_column_name(
                         key
                     )
-                    expr_id = right_input[snowpark_name]._expression.expr_id
+                    expr_id = right_container.dataframe[
+                        snowpark_name
+                    ]._expression.expr_id
                     updated_key = COLUMN_METADATA_COLLISION_KEY.format(
                         expr_id=expr_id, key=snowpark_name
                     )
@@ -498,49 +651,137 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
                 except Exception:
                     # ignore any errors that happens while fetching the metadata
                     pass
+    return column_metadata
+
+
+def _build_joined_schema(
+    snowpark_columns: list[str],
+    left_input: DataFrame,
+    right_input: DataFrame,
+    outer_join_columns: Optional[list[ColumnNames]] = None,
+) -> Callable[[], StructType]:
+    """
+    Builds a lazy schema for the joined dataframe, based on the given snowpark_columns and input dataframes.
+    In case of full outer joins, we need a separate target_snowpark_columns, since join columns will have different
+    names in the output than in any input.
+    """
+
+    def _schema_getter() -> StructType:
+        all_fields = left_input.schema.fields + right_input.schema.fields
+        fields: dict[str, StructField] = {f.name: f for f in all_fields}
+
+        if outer_join_columns:
+            visible_columns = [c for c in outer_join_columns if not c.is_hidden]
+            assert len(snowpark_columns) == len(visible_columns)
+
+            result_fields = []
+            visible_idx = 0
+            for col in outer_join_columns:
+                if col.is_hidden:
+                    source_field = fields[col.snowpark_name]
+                    result_fields.append(
+                        StructField(
+                            col.snowpark_name,
+                            source_field.datatype,
+                            source_field.nullable,
+                        )
+                    )
+                else:
+                    source_field = fields[snowpark_columns[visible_idx]]
+                    result_fields.append(
+                        StructField(
+                            col.snowpark_name,
+                            source_field.datatype,
+                            source_field.nullable,
+                        )
+                    )
+                    visible_idx += 1
 
-    result_container = DataFrameContainer.create_with_column_mapping(
-        dataframe=result,
-        spark_column_names=spark_cols_after_join,
-        snowpark_column_names=snowpark_cols_after_join_deduplicated,
-        column_metadata=column_metadata,
-        column_qualifiers=qualifiers,
-        hidden_columns=hidden_columns,
-        snowpark_column_types=snowpark_col_types,
-    )
-
-    if rel.join.using_columns:
-        # When join 'using_columns', the 'join columns' should go first in result DF.
-        idxs_to_shift = [
-            spark_cols_after_join.index(left_col_name)
-            for left_col_name in case_corrected_left_columns
-        ]
-
-        def reorder(lst: list) -> list:
-            to_move = [lst[i] for i in idxs_to_shift]
-            remaining = [el for i, el in enumerate(lst) if i not in idxs_to_shift]
-            return to_move + remaining
+            return StructType(result_fields)
 
-        # Create reordered DataFrame
-        reordered_df = result_container.dataframe.select(
-            [snowpark_fn.col(c) for c in reorder(result_container.dataframe.columns)]
+        return StructType(
+            [
+                StructField(name, fields[name].datatype, fields[name].nullable)
+                for name in snowpark_columns
+            ]
         )
 
-        # Create new container with reordered metadata
-        original_df = result_container.dataframe
-        return DataFrameContainer.create_with_column_mapping(
-            dataframe=reordered_df,
-            spark_column_names=reorder(result_container.column_map.get_spark_columns()),
-            snowpark_column_names=reorder(
-                result_container.column_map.get_snowpark_columns()
-            ),
-            column_metadata=column_metadata,
-            column_qualifiers=reorder(qualifiers),
-            table_name=result_container.table_name,
-            cached_schema_getter=lambda: snowpark.types.StructType(
-                reorder(original_df.schema.fields)
-            ),
-            hidden_columns=hidden_columns,
+    return _schema_getter
+
+
+def _make_struct_column(
+    container: DataFrameContainer, snowpark_name: str
+) -> tuple[snowpark.Column, StructType]:
+    column_metadata: dict = {}
+    for c in container.column_map.columns:
+        column_metadata[c.snowpark_name] = c
+
+    args: list[Column] = []
+    struct_fields: list[StructField] = []
+    for f in container.dataframe.schema.fields:
+        c = column_metadata[f.name]
+        if c.is_hidden:
+            continue
+        args.append(snowpark_fn.lit(c.spark_name))
+        args.append(snowpark_fn.col(c.snowpark_name))
+        struct_fields.append(
+            StructField(c.spark_name, f.datatype, f.nullable, _is_column=False)
         )
 
-    return result_container
+    struct_type = StructType(struct_fields, structured=True)
+    struct_col: snowpark.Column = (
+        snowpark_fn.object_construct_keep_null(*args)
+        .cast(struct_type)
+        .alias(snowpark_name)
+    )
+    return struct_col, struct_type
+
+
+def _construct_join_with_column(
+    container: DataFrameContainer, snowpark_name: str, is_struct: bool
+) -> tuple[Column, StructType]:
+    if is_struct:
+        return _make_struct_column(container, snowpark_name)
+    else:
+        # the dataframe must have a single field
+        cols = [
+            c.snowpark_name for c in container.column_map.columns if not c.is_hidden
+        ]
+        assert (
+            len(cols) == 1
+        ), "A non-struct dataframe must have a single column in joinWith"
+        field = None
+        for f in container.dataframe.schema.fields:
+            if f.name == cols[0]:
+                field = f
+                break
+        assert field is not None
+        col = snowpark_fn.col(field.name).alias(snowpark_name)
+        col_type = field.datatype
+        return col, col_type
+
+
+def _join_with_nullability(join_type: str) -> tuple[bool, bool]:
+    """
+    Returns the nullability for the left and right result columns of a joinWith operation.
+
+    The tuple corresponds to (left_nullable, right_nullable) and depends on the join type:
+    - "inner" or "cross": both columns are non-nullable
+    - "left": left is non-nullable, right is nullable
+    - "right": left is nullable, right is non-nullable
+    - "full_outer": both columns are nullable
+
+    Raises:
+        IllegalArgumentException: If the provided join type is unsupported.
+    """
+    match join_type:
+        case "inner" | "cross":
+            return False, False
+        case "left":
+            return False, True
+        case "right":
+            return True, False
+        case "full_outer":
+            return True, True
+        case _:
+            raise IllegalArgumentException(f"Unsupported join type '{join_type}'.")