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

snowflake/snowpark_connect/relation/map_join.py

@@ -2,6 +2,7 @@
 # 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
@@ -11,12 +12,12 @@ from pyspark.errors import AnalysisException
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
+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.column_qualifier import ColumnQualifier
 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
@@ -60,9 +61,7 @@ class JoinInfo:
     join_type: str
     condition_type: ConditionType
     join_columns: Optional[list[str]]
-
-    def has_join_condition(self) -> bool:
-        return self.condition_type == ConditionType.JOIN_CONDITION
+    just_left_columns: bool
 
     def is_using_columns(self):
         return self.condition_type == ConditionType.USING_COLUMNS
@@ -77,196 +76,26 @@ 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)
-    join_type = join_info.join_type
-
-    if join_info.has_join_condition():
-        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),
-            )
-        result: snowpark.DataFrame = left_input.join(
-            right=right_input,
-            on=join_expression.col,
-            how="inner" if join_info.join_type == "cross" else join_info.join_type,
-            lsuffix="_left",
-            rsuffix="_right",
-        )
-    elif join_info.is_using_columns():
-        # TODO: disambiguate snowpark columns for all join condition types
-        # disambiguation temporarily done only for using_columns/natural joins to reduce changes
-        left_container, right_container = _disambiguate_snowpark_columns(
-            left_container, right_container
-        )
-        left_input = left_container.dataframe
-        right_input = right_container.dataframe
-
-        join_columns = join_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")
-
-        # 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(
-                        spark_name, return_first=True
-                    )
-                ],
-                right_input[
-                    right_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                        spark_name, return_first=True
-                    )
-                ],
+    match join_info.condition_type:
+        case ConditionType.JOIN_CONDITION:
+            result_container = _join_using_condition(
+                left_container, right_container, join_info, rel
             )
-            for spark_name in join_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])
-
-            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]
-            ]
-            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
-        # For outer joins without a condition, we need to use a TRUE condition
-        # to match Spark's behavior.
-        if join_type in ["left", "right", "full_outer"]:
-            result: snowpark.DataFrame = left_input.join(
-                right=right_input,
-                on=snowpark_fn.lit(True),
-                how=join_type,
-            )
-        else:
-            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:
-        if not join_info.is_using_columns():
-            spark_cols_after_join: list[str] = (
-                left_container.column_map.get_spark_columns()
-                + right_container.column_map.get_spark_columns()
+        case ConditionType.USING_COLUMNS:
+            result_container = _join_using_columns(
+                left_container, right_container, join_info
             )
-            qualifiers: list[set[ColumnQualifier]] = (
-                left_container.column_map.get_qualifiers()
-                + right_container.column_map.get_qualifiers()
+        case _:
+            result_container = _join_unconditionally(
+                left_container, right_container, join_info
             )
-        else:
-            # get columns after join
-            joined_columns = left_container.column_map.get_columns_after_join(
-                right_container.column_map, join_info.join_columns
-            )
-            spark_cols_after_join: list[str] = [c.spark_name for c in joined_columns]
-            qualifiers: list[set[ColumnQualifier]] = [
-                c.qualifiers for c in joined_columns
-            ]
-
-    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:
-                column_metadata[key] = value
-            else:
-                # In case of collision, use snowpark's column's expr_id as prefix.
-                # this is a temporary solution until SNOW-1926440 is resolved.
-                try:
-                    snowpark_name = right_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                        key
-                    )
-                    expr_id = right_input[snowpark_name]._expression.expr_id
-                    updated_key = COLUMN_METADATA_COLLISION_KEY.format(
-                        expr_id=expr_id, key=snowpark_name
-                    )
-                    column_metadata[updated_key] = value
-                except Exception:
-                    # ignore any errors that happens while fetching the metadata
-                    pass
-
-    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
@@ -283,47 +112,266 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
     # since both columns are replaced with a coalesced column
     if (
         join_info.is_using_columns()
-        and join_type == "full_outer"
+        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)
 
-    if join_info.is_using_columns():
-        # When join 'using_columns', the 'join columns' should go first in result DF.
-        # we're only shifting left side columns, since we dropped the right-side ones
-        idxs_to_shift = left_container.column_map.get_column_indexes(
-            join_info.join_columns
+    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) :]
 
-        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
+        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]
 
-        # Create reordered DataFrame
-        reordered_df = result_container.dataframe.select(
-            [snowpark_fn.col(c) for c in reorder(result_container.dataframe.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,
+            ),
         )
 
-        # Create new container with reordered metadata
-        original_df = result_container.dataframe
+    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=reordered_df,
-            spark_column_names=reorder(result_container.column_map.get_spark_columns()),
-            snowpark_column_names=reorder(
-                result_container.column_map.get_snowpark_columns()
+            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
             ),
-            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)
+        )
+
+    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),
         )
 
-    return result_container
+    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(
@@ -372,6 +420,10 @@ def _get_join_info(
     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
 
@@ -381,11 +433,17 @@ def _get_join_info(
     elif is_using_columns:
         condition_type = ConditionType.USING_COLUMNS
 
-    return JoinInfo(join_type, condition_type, join_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"]
+
+    return JoinInfo(join_type, condition_type, join_columns, just_left_columns)
 
 
 def _disambiguate_snowpark_columns(
-    left: DataFrameContainer, right: DataFrameContainer
+    left: DataFrameContainer, right: DataFrameContainer, rel: relation_proto.Relation
 ) -> tuple[DataFrameContainer, DataFrameContainer]:
     conflicting_snowpark_columns = left.column_map.get_conflicting_snowpark_columns(
         right.column_map
@@ -394,14 +452,24 @@ def _disambiguate_snowpark_columns(
     if not conflicting_snowpark_columns:
         return left, right
 
-    # rename and create new containers
-    return _disambiguate_container(
-        left, conflicting_snowpark_columns
-    ), _disambiguate_container(right, conflicting_snowpark_columns)
+    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]
+    container: DataFrameContainer,
+    conflicting_snowpark_columns: set[str],
+    plan_id: Optional[int],
 ) -> DataFrameContainer:
     column_map = container.column_map
     disambiguated_columns = []
@@ -420,25 +488,87 @@ def _disambiguate_container(
 
     disambiguated_df = container.dataframe.select(*disambiguated_columns)
 
-    def _get_new_schema():
-        old_schema = container.dataframe.schema
-        if not old_schema.fields:
-            return StructType([])
-
-        new_fields = []
-        for i, name in enumerate(disambiguated_snowpark_names):
-            f = old_schema.fields[i]
-            new_fields.append(
-                StructField(name, f.datatype, nullable=f.nullable, _is_column=True)
-            )
-        return StructType(new_fields)
+    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)
+            ]
+        )
 
-    return DataFrameContainer.create_with_column_mapping(
+    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=_get_new_schema,
+        cached_schema_getter=_schema_getter,
     )
+
+    # 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
+
+
+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:
+                column_metadata[key] = value
+            else:
+                # In case of collision, use snowpark's column's expr_id as prefix.
+                # this is a temporary solution until SNOW-1926440 is resolved.
+                try:
+                    snowpark_name = right_container.column_map.get_snowpark_column_name_from_spark_column_name(
+                        key
+                    )
+                    expr_id = right_container.dataframe[
+                        snowpark_name
+                    ]._expression.expr_id
+                    updated_key = COLUMN_METADATA_COLLISION_KEY.format(
+                        expr_id=expr_id, key=snowpark_name
+                    )
+                    column_metadata[updated_key] = value
+                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,
+    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.
+    """
+
+    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
+
+        assert len(snowpark_columns) == len(target_names)
+
+        return StructType(
+            [
+                StructField(
+                    target_names[i], fields[name].datatype, fields[name].nullable
+                )
+                for i, name in enumerate(snowpark_columns)
+            ]
+        )
+
+    return _schema_getter