snowpark-connect 0.32.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

snowflake/snowpark_connect/relation/map_join.py

@@ -1,16 +1,23 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
+import dataclasses
+from collections.abc import Callable
+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 import AnalysisException
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import JoinColumnNameMap
-from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
+from snowflake.snowpark import DataFrame
+from snowflake.snowpark.types import StructField, StructType
+from snowflake.snowpark_connect.column_name_handler import (
+    JoinColumnNameMap,
+    make_unique_snowpark_name,
+)
 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
@@ -43,6 +50,23 @@ 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}"
 
 
+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
+
+    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)
@@ -52,18 +76,321 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
     left_container = filter_metadata_columns(left_container)
     right_container = filter_metadata_columns(right_container)
 
-    left_input: snowpark.DataFrame = left_container.dataframe
-    right_input: snowpark.DataFrame = right_container.dataframe
+    # if there are any conflicting snowpark columns, this is the time to rename them
+    left_container, right_container = _disambiguate_snowpark_columns(
+        left_container, right_container, rel
+    )
+
+    join_info = _get_join_info(rel, left_container, right_container)
+
+    match join_info.condition_type:
+        case ConditionType.JOIN_CONDITION:
+            result_container = _join_using_condition(
+                left_container, right_container, join_info, rel
+            )
+        case ConditionType.USING_COLUMNS:
+            result_container = _join_using_columns(
+                left_container, right_container, join_info
+            )
+        case _:
+            result_container = _join_unconditionally(
+                left_container, right_container, join_info
+            )
+
+    # Fix for USING join column references with different plan IDs
+    # After a USING join, references to the right dataframe's columns should resolve
+    # to the result dataframe that contains the merged columns
+    if (
+        join_info.is_using_columns()
+        and rel.join.right.HasField("common")
+        and rel.join.right.common.HasField("plan_id")
+    ):
+        right_plan_id = rel.join.right.common.plan_id
+        set_plan_id_map(right_plan_id, result_container)
+
+    # For FULL OUTER joins, we also need to map the left dataframe's plan_id
+    # since both columns are replaced with a coalesced column
+    if (
+        join_info.is_using_columns()
+        and join_info.join_type == "full_outer"
+        and rel.join.left.HasField("common")
+        and rel.join.left.common.HasField("plan_id")
+    ):
+        left_plan_id = rel.join.left.common.plan_id
+        set_plan_id_map(left_plan_id, result_container)
+
+    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
+        ),
+    )
+
+
+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
+        ):
+            exception = AnalysisException(
+                USING_COLUMN_NOT_FOUND_ERROR.format(
+                    column, side, container.column_map.get_spark_columns()
+                )
+            )
+            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,
+    )
+
+    # 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
+    )
+
+    # For outer joins, we need to preserve join keys from both sides using COALESCE
+    if info.join_type == "full_outer":
+        coalesced_columns = []
+        coalesced_column_names = []
+        for i, (left_col, right_col) in enumerate(snowpark_using_columns):
+            # spark uses the left side spark name
+            spark_name = columns[i].spark_name
+            new_snowpark_name = make_unique_snowpark_name(spark_name)
+            coalesced_col = snowpark_fn.coalesce(left_col, right_col).alias(
+                new_snowpark_name
+            )
+            coalesced_columns.append(coalesced_col)
+            coalesced_column_names.append((spark_name, new_snowpark_name))
+
+        # join columns need to be replaced, so we need the original names for schema lookup later
+        snowpark_names_for_schema_lookup = [c.snowpark_name for c in columns]
+
+        # we need to use the coalesced columns instead of the left-side join columns
+        columns = columns[len(join_columns) :]
+
+        non_join_columns = [snowpark_fn.col(c.snowpark_name) for c in columns]
+        result = joined_df.select(coalesced_columns + non_join_columns)
+
+        spark_names = [spark_name for spark_name, _ in coalesced_column_names] + [
+            c.spark_name for c in columns
+        ]
+        snowpark_names = [
+            snowpark_name for _, snowpark_name in coalesced_column_names
+        ] + [c.snowpark_name for c in columns]
+        qualifiers = ([set()] * len(join_columns)) + [c.qualifiers for c in columns]
+
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result,
+            spark_column_names=spark_names,
+            snowpark_column_names=snowpark_names,
+            column_metadata=_combine_metadata(left_container, right_container),
+            column_qualifiers=qualifiers,
+            cached_schema_getter=_build_joined_schema(
+                snowpark_names_for_schema_lookup,
+                left_input,
+                right_input,
+                snowpark_names,
+            ),
+        )
+
+    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
+            ),
+        )
+
+    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
+        ),
+    )
+
+
+def _join_using_condition(
+    left_container: DataFrameContainer,
+    right_container: DataFrameContainer,
+    info: JoinInfo,
+    rel: relation_proto.Relation,
+) -> 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,
+                right_container.column_map,
+            ),
+            typer=JoinExpressionTyper(left_input, right_input),
+        )
+
+    result: snowpark.DataFrame = left_input.join(
+        right=right_input,
+        on=join_expression.col,
+        how=info.join_type,
+    )
+
+    # 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)
+
+    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
+
+    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
+        ),
+    )
+
+
+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
-    using_columns = rel.join.using_columns
+    join_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()
+        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
         ]
-        using_columns = common_spark_columns
+        join_columns = common_spark_columns
 
     match rel.join.join_type:
         case relation_proto.Join.JOIN_TYPE_UNSPECIFIED:
@@ -90,202 +417,108 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
             attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
             raise exception
 
-    # This handles case sensitivity for using_columns
-    case_corrected_right_columns: list[str] = []
+    has_join_condition = rel.join.HasField("join_condition")
+    is_using_columns = bool(join_columns)
 
-    if rel.join.HasField("join_condition"):
-        assert not using_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"
 
-        left_columns = list(left_container.column_map.spark_to_col.keys())
-        right_columns = list(right_container.column_map.spark_to_col.keys())
+    if has_join_condition:
+        assert not is_using_columns
 
-        # 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),
-            )
-        result: snowpark.DataFrame = left_input.join(
-            right=right_input,
-            on=join_expression.col,
-            how=join_type,
-            lsuffix="_left",
-            rsuffix="_right",
-        )
-    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
-            )
-            is None
-            for c in using_columns
-        ):
-            exception = 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(),
-                )
-            )
-            attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
-            raise exception
-        if any(
-            right_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                c, allow_non_exists=True, return_first=True
-            )
-            is None
-            for c in using_columns
-        ):
-            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(),
-                )
-            )
-            attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
-            raise exception
+    condition_type = ConditionType.NO_CONDITION
+    if has_join_condition:
+        condition_type = ConditionType.JOIN_CONDITION
+    elif is_using_columns:
+        condition_type = ConditionType.USING_COLUMNS
 
-        # 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(
-            left_container.column_map.get_snowpark_column_names_from_spark_column_names(
-                list(using_columns), 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
-            )
-        )
-        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
-                    )
-                ],
+    # 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"]
+
+    return JoinInfo(join_type, condition_type, join_columns, just_left_columns)
+
+
+def _disambiguate_snowpark_columns(
+    left: DataFrameContainer, right: DataFrameContainer, rel: relation_proto.Relation
+) -> tuple[DataFrameContainer, DataFrameContainer]:
+    conflicting_snowpark_columns = left.column_map.get_conflicting_snowpark_columns(
+        right.column_map
+    )
+
+    if not conflicting_snowpark_columns:
+        return left, right
+
+    left_plan = rel.join.left.common.plan_id
+    right_plan = rel.join.right.common.plan_id
+
+    if left_plan == right_plan:
+        # don't overwrite plan_id map for self joins
+        right_plan = None
+
+    # rename and create new right container
+    # TODO: rename both sides after SNOW-2382499
+    return left, _disambiguate_container(
+        right, conflicting_snowpark_columns, right_plan
+    )
+
+
+def _disambiguate_container(
+    container: DataFrameContainer,
+    conflicting_snowpark_columns: set[str],
+    plan_id: Optional[int],
+) -> DataFrameContainer:
+    column_map = container.column_map
+    disambiguated_columns = []
+    disambiguated_snowpark_names = []
+    for c in column_map.columns:
+        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)
             )
-            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,
-        )
-        # For outer joins, we need to preserve join keys from both sides using COALESCE
-        if join_type == "full_outer":
-            coalesced_columns = []
-            columns_to_drop = []
-            for i, (left_col, right_col) in enumerate(snowpark_using_columns):
-                # Use the original user-specified column name to preserve case sensitivity
-                original_column_name = rel.join.using_columns[i]
-                coalesced_col = snowpark_fn.coalesce(left_col, right_col).alias(
-                    original_column_name
-                )
-                coalesced_columns.append(coalesced_col)
-                columns_to_drop.extend([left_col, right_col])
+        else:
+            disambiguated_snowpark_names.append(c.snowpark_name)
+            disambiguated_columns.append(snowpark_fn.col(c.snowpark_name))
+
+    disambiguated_df = container.dataframe.select(*disambiguated_columns)
 
-            other_columns = [
-                snowpark_fn.col(col_name)
-                for col_name in joined_df.columns
-                if col_name not in [col.getName() for col in columns_to_drop]
+    def _schema_getter() -> StructType:
+        fields = container.dataframe.schema.fields
+        return StructType(
+            [
+                StructField(name, fields[i].datatype, fields[i].nullable)
+                for i, name in enumerate(disambiguated_snowpark_names)
             ]
-            result = joined_df.select(coalesced_columns + other_columns)
-        else:
-            result = joined_df.drop(*(right for _, right in snowpark_using_columns))
-    else:
-        if 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
-        result: snowpark.DataFrame = left_input.join(
-            right=right_input,
-            how=join_type,
         )
 
-    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: list[str] = left_container.column_map.get_spark_columns()
-        qualifiers = left_container.column_map.get_qualifiers()
-    else:
-        # Add Spark columns and plan_ids from left DF
-        spark_cols_after_join: list[str] = list(
-            left_container.column_map.get_spark_columns()
-        ) + [
-            spark_col
-            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
-            ]  # this is to make sure we only remove the column once
-        ]
+    disambiguated_container = 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=container.table_name,
+        cached_schema_getter=_schema_getter,
+    )
 
-        qualifiers: list[set[ColumnQualifier]] = list(
-            left_container.column_map.get_qualifiers()
-        ) + [
-            {right_container.column_map.get_qualifier_for_spark_column(spark_col)}
-            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
-            ]  # this is to make sure we only remove the column once]
-        ]
+    # 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 plan_id is not None:
+        set_plan_id_map(plan_id, disambiguated_container)
+
+    return disambiguated_container
 
-    column_metadata = {}
-    if left_container.column_map.column_metadata:
-        column_metadata.update(left_container.column_map.column_metadata)
 
+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:
@@ -297,7 +530,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
                     )
@@ -305,68 +540,35 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
                 except Exception:
                     # ignore any errors that happens while fetching the metadata
                     pass
+    return column_metadata
 
-    result_container = DataFrameContainer.create_with_column_mapping(
-        dataframe=result,
-        spark_column_names=spark_cols_after_join,
-        snowpark_column_names=result.columns,
-        column_metadata=column_metadata,
-        column_qualifiers=qualifiers,
-    )
 
-    # Fix for USING join column references with different plan IDs
-    # After a USING join, references to the right dataframe's columns should resolve
-    # to the result dataframe that contains the merged columns
-    if (
-        using_columns
-        and rel.join.right.HasField("common")
-        and rel.join.right.common.HasField("plan_id")
-    ):
-        right_plan_id = rel.join.right.common.plan_id
-        set_plan_id_map(right_plan_id, result_container)
+def _build_joined_schema(
+    snowpark_columns: list[str],
+    left_input: DataFrame,
+    right_input: DataFrame,
+    target_snowpark_columns: Optional[list[str]] = 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.
+    """
 
-    # For FULL OUTER joins, we also need to map the left dataframe's plan_id
-    # since both columns are replaced with a coalesced column
-    if (
-        using_columns
-        and join_type == "full_outer"
-        and rel.join.left.HasField("common")
-        and rel.join.left.common.HasField("plan_id")
-    ):
-        left_plan_id = rel.join.left.common.plan_id
-        set_plan_id_map(left_plan_id, result_container)
-
-    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 _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}
+        target_names = target_snowpark_columns or snowpark_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
+        assert len(snowpark_columns) == len(target_names)
 
-        # Create reordered DataFrame
-        reordered_df = result_container.dataframe.select(
-            [snowpark_fn.col(c) for c in reorder(result_container.dataframe.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)
-            ),
+        return StructType(
+            [
+                StructField(
+                    target_names[i], fields[name].datatype, fields[name].nullable
+                )
+                for i, name in enumerate(snowpark_columns)
+            ]
         )
 
-    return result_container
+    return _schema_getter