Remove unecessary dropnas (#1292)

snowflakedb · Mar 7, 2024 · 58c0acb · 58c0acb
1 parent 3ce2356
commit 58c0acb
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,10 @@
   - Session.use_database
   - Session.use_role
 
+### Bug Fixes
+
+- Fixed a bug in Local Testing's implementation of LEFT ANTI and LEFT SEMI joins where rows with null values are dropped.
+
 ### Deprecations:
 
 - Deprecated `Session.get_fully_qualified_current_schema`. Consider using `Session.get_fully_qualified_name_if_possible` instead.

diff --git a/src/snowflake/snowpark/mock/_plan.py b/src/snowflake/snowpark/mock/_plan.py
@@ -854,9 +854,9 @@ def outer_join(base_df):
             )
             sf_types = result_df.sf_types
             if "SEMI" in source_plan.join_type.sql:  # left semi
-                result_df = left[outer_join(left)].dropna()
+                result_df = left[outer_join(left)]
             elif "ANTI" in source_plan.join_type.sql:  # left anti
-                result_df = left[~outer_join(left)].dropna()
+                result_df = left[~outer_join(left)]
             elif "LEFT" in source_plan.join_type.sql:  # left outer join
                 # rows from LEFT that did not get matched
                 unmatched_left = left[~outer_join(left)]

diff --git a/tests/integ/test_dataframe.py b/tests/integ/test_dataframe.py
@@ -1193,6 +1193,12 @@ def test_join_left_anti(session):
     expected = [Row(3, 3), Row(4, 4)]
     assert sorted(res, key=lambda r: r[0]) == expected
 
+    df3 = session.create_dataframe(
+        [[i if i & 2 else None] for i in range(3, 8)], schema=["id"]
+    )
+    res = df3.join(df1, "id", "leftanti").collect()
+    assert res == [Row(ID=None), Row(ID=None)]
+
 
 @pytest.mark.localtest
 def test_join_left_outer(session):