From 2901895568e9b71905ddd669dc08d86336211f57 Mon Sep 17 00:00:00 2001
From: azhan <andong.zhan@snowflake.com>
Date: Wed, 28 Aug 2024 09:20:37 -0700
Subject: [PATCH] resolve comments and fix tests

---
 src/snowflake/snowpark/modin/plugin/_internal/join_utils.py | 2 +-
 .../modin/plugin/compiler/snowflake_query_compiler.py       | 2 +-
 tests/integ/modin/frame/test_isin.py                        | 2 +-
 tests/integ/modin/frame/test_iterrows.py                    | 2 +-
 tests/integ/modin/series/test_shift.py                      | 6 +-----
 5 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/src/snowflake/snowpark/modin/plugin/_internal/join_utils.py b/src/snowflake/snowpark/modin/plugin/_internal/join_utils.py
index be36b004ed4..846f3c64079 100644
--- a/src/snowflake/snowpark/modin/plugin/_internal/join_utils.py
+++ b/src/snowflake/snowpark/modin/plugin/_internal/join_utils.py
@@ -173,7 +173,7 @@ def join(
     ), f"Invalid join type: {how}. Allowed values are {get_args(JoinTypeLit)}"
 
     def assert_snowpark_pandas_types_match() -> None:
-        """If Snowpark pandas types does not match, then a ValueError will be raised."""
+        """If Snowpark pandas types do not match, then a ValueError will be raised."""
         left_types = [
             left.snowflake_quoted_identifier_to_snowpark_pandas_type.get(id, None)
             for id in left_on
diff --git a/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py b/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
index 795e6bd6e76..126c2c5c324 100644
--- a/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
+++ b/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
@@ -1499,7 +1499,7 @@ def _shift_values_axis_0(
         row_position_quoted_identifier = frame.row_position_snowflake_quoted_identifier
 
         fill_value_dtype = infer_object_type(fill_value)
-        fill_value = pandas_lit(fill_value) if fill_value is not None else None
+        fill_value = None if pd.isna(fill_value) else pandas_lit(fill_value)
 
         def shift_expression_and_type(
             quoted_identifier: str, dtype: DataType
diff --git a/tests/integ/modin/frame/test_isin.py b/tests/integ/modin/frame/test_isin.py
index cc6113c7466..cd560a5715a 100644
--- a/tests/integ/modin/frame/test_isin.py
+++ b/tests/integ/modin/frame/test_isin.py
@@ -260,7 +260,7 @@ def test_isin_dataframe_values_type_negative():
 )
 def test_isin_timedelta(values):
     native_df = native_pd.DataFrame({"a": [1, 2, 3], "b": [None, 4, 2]}).astype(
-        "timedelta64[ns]"
+        {"b": "timedelta64[ns]"}
     )
     snow_df = pd.DataFrame(native_df)
 
diff --git a/tests/integ/modin/frame/test_iterrows.py b/tests/integ/modin/frame/test_iterrows.py
index 843848c6a5c..fc415b2daf5 100644
--- a/tests/integ/modin/frame/test_iterrows.py
+++ b/tests/integ/modin/frame/test_iterrows.py
@@ -53,7 +53,7 @@ def assert_iterators_equal(snowpark_iterator, native_iterator):
         ),
         # empty df
         native_pd.DataFrame([]),
-        native_pd.DataFrame({"ts": native_pd.timedelta_range(10, periods=10)}),
+        native_pd.DataFrame({"ts": native_pd.timedelta_range(10, periods=4)}),
     ],
 )
 def test_df_iterrows(native_df):
diff --git a/tests/integ/modin/series/test_shift.py b/tests/integ/modin/series/test_shift.py
index 7f27c4d313b..f5d4169026e 100644
--- a/tests/integ/modin/series/test_shift.py
+++ b/tests/integ/modin/series/test_shift.py
@@ -46,11 +46,7 @@ def test_series_with_values_shift(series, periods, fill_value):
         lambda s: s.shift(
             periods=periods,
             fill_value=pd.Timedelta(fill_value)
-            if isinstance(
-                s, native_pd.Series
-            )  # pandas does not support fill int to timedelta
-            and s.dtype == "timedelta64[ns]"
-            and fill_value is not no_default
+            if s.dtype == "timedelta64[ns]" and fill_value is not no_default
             else fill_value,
         ),
     )