[Local testing] Fixes various local testing bugs

CHANGELOG.md

@@ -55,6 +55,17 @@
 - Fixed a bug where transposing a dataframe would map `Timedelta` index levels to integer column levels.
 - Fixed a bug where `Resampler` methods on timedelta columns would produce integer results.
 
+### Snowpark Local Testing Updates
+
+#### Bug Fixes
+
+- Fixed a bug where nullable columns were annotated wrongly.
+- Fixed a bug where the `date_add` and `date_sub` functions failed for `NULL` values.
+- Fixed a bug where `equal_null` could fail inside a merge statement.
+- Fixed a bug where `row_number` could fail inside a Window function.
+- Fixed a bug where updates could fail when the source is the result of a join.
+
+
 ## 1.22.1 (2024-09-11)
 This is a re-release of 1.22.0. Please refer to the 1.22.0 release notes for detailed release content.
 

src/snowflake/snowpark/mock/_functions.py

@@ -1530,6 +1530,8 @@ def add_months(scalar, date, duration):
 
 
 def add_timedelta(unit, date, duration, scalar=1):
+    if date is None:
+        return date
     return date + datetime.timedelta(**{f"{unit}s": float(duration) * scalar})
 
 

src/snowflake/snowpark/mock/_plan.py

@@ -640,9 +640,6 @@ def execute_mock_plan(
                     if isinstance(exp.child, Attribute):
                         quoted_name = quote_name(exp.name)
                         expr_to_alias[exp.child.expr_id] = quoted_name
-                        for k, v in expr_to_alias.items():
-                            if v == exp.child.name:
-                                expr_to_alias[k] = quoted_name
 
         df = pd.concat(data, axis=1)
         result_df = TableEmulator(
@@ -1289,7 +1286,7 @@ def outer_join(base_df):
             join_condition = calculate_expression(
                 source_plan.join_expr, cartesian_product, analyzer, expr_to_alias
             )
-            join_result = cartesian_product[join_condition]
+            join_result = cartesian_product[join_condition].reset_index(drop=True)
             join_result.sf_types = cartesian_product.sf_types
 
             # TODO [GA]: # ERROR_ON_NONDETERMINISTIC_MERGE is by default True, raise error if
@@ -1711,7 +1708,7 @@ def calculate_expression(
                 exp.datatype = StringType(len(exp.value))
             res = ColumnEmulator(
                 data=[exp.value for _ in range(len(input_data))],
-                sf_type=ColumnType(exp.datatype, False),
+                sf_type=ColumnType(exp.datatype, nullable=exp.value is None),
                 dtype=object,
             )
             res.index = input_data.index
@@ -2062,6 +2059,17 @@ def _match_pattern(row) -> bool:
                 isinstance(lower, UnboundedPreceding),
                 isinstance(upper, UnboundedFollowing),
             )
+
+        # Reorder windows to match the index order in res_index
+        reordered_windows = []
+        for idx in res_index:
+            for w in windows:
+                if idx in w.index:
+                    reordered_windows.append(w)
+                    break
+
+        windows = reordered_windows
+
         # compute window function:
         if isinstance(window_function, (FunctionExpression,)):
             res_cols = []