SNOW-1662275: Add support for left_index and right_index for pd.merge…

…_asof Signed-off-by: Naren Krishna <[email protected]>
snowflakedb · Sep 17, 2024 · 28cc861 · 28cc861
1 parent e93cd68
commit 28cc861
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,7 +12,7 @@
 
 - Added support for `TimedeltaIndex.mean` method.
 - Added support for some cases of aggregating `Timedelta` columns on `axis=0` with `agg` or `aggregate`.
-- Added support for `by`, `left_by`, and `right_by` for `pd.merge_asof`.
+- Added support for `by`, `left_by`, `right_by`, `left_index`, and `right_index` for `pd.merge_asof`.
 
 #### Bug Fixes
 

diff --git a/docs/source/modin/supported/general_supported.rst b/docs/source/modin/supported/general_supported.rst
@@ -38,8 +38,7 @@ Data manipulations
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``merge``                   | P                               | ``validate``                     | ``N`` if param ``validate`` is given               |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
-| ``merge_asof``              | P                               | ``left_index``, ``right_index``, | ``N`` if param ``direction`` is ``nearest``.       |
-|                             |                                 | , ``suffixes``, ``tolerance``    |                                                    |
+| ``merge_asof``              | P                               | ``suffixes``, ``tolerance``      | ``N`` if param ``direction`` is ``nearest``        |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``merge_ordered``           | N                               |                                  |                                                    |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+

diff --git a/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py b/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
@@ -7381,10 +7381,10 @@ def merge_asof(
         SnowflakeQueryCompiler
         """
         # TODO: SNOW-1634547: Implement remaining parameters by leveraging `merge` implementation
-        if left_index or right_index or tolerance or suffixes != ("_x", "_y"):
+        if tolerance or suffixes != ("_x", "_y"):
             ErrorMessage.not_implemented(
                 "Snowpark pandas merge_asof method does not currently support parameters "
-                + "'left_index', 'right_index', 'suffixes', or 'tolerance'"
+                + "'suffixes', or 'tolerance'"
             )
         if direction not in ("backward", "forward"):
             ErrorMessage.not_implemented(

diff --git a/tests/integ/modin/test_merge_asof.py b/tests/integ/modin/test_merge_asof.py
@@ -231,6 +231,37 @@ def test_merge_asof_left_right_on(
     assert_snowpark_pandas_equal_to_pandas(snow_output, native_output)
 
 
+@allow_exact_matches
+@direction
+@sql_count_checker(query_count=1, join_count=1)
+def test_merge_asof_left_right_index(allow_exact_matches, direction):
+    native_left = native_pd.DataFrame({"left_val": ["a", "b", "c"]}, index=[1, 5, 10])
+    native_right = native_pd.DataFrame(
+        {"right_val": [1, 2, 3, 6, 7]}, index=[1, 2, 3, 6, 7]
+    )
+
+    snow_left = pd.DataFrame(native_left)
+    snow_right = pd.DataFrame(native_right)
+
+    native_output = native_pd.merge_asof(
+        native_left,
+        native_right,
+        left_index=True,
+        right_index=True,
+        direction=direction,
+        allow_exact_matches=allow_exact_matches,
+    )
+    snow_output = pd.merge_asof(
+        snow_left,
+        snow_right,
+        left_index=True,
+        right_index=True,
+        direction=direction,
+        allow_exact_matches=allow_exact_matches,
+    )
+    assert_snowpark_pandas_equal_to_pandas(snow_output, native_output)
+
+
 @pytest.mark.parametrize("by", ["ticker", ["ticker"]])
 @sql_count_checker(query_count=1, join_count=1)
 def test_merge_asof_by(left_right_timestamp_data, by):
@@ -399,15 +430,7 @@ def test_merge_asof_params_unsupported(left_right_timestamp_data):
         NotImplementedError,
         match=(
             "Snowpark pandas merge_asof method does not currently support parameters "
-            + "'left_index', 'right_index', 'suffixes', or 'tolerance'"
-        ),
-    ):
-        pd.merge_asof(left_snow_df, right_snow_df, left_index=True, right_index=True)
-    with pytest.raises(
-        NotImplementedError,
-        match=(
-            "Snowpark pandas merge_asof method does not currently support parameters "
-            + "'left_index', 'right_index', 'suffixes', or 'tolerance'"
+            + "'suffixes', or 'tolerance'"
         ),
     ):
         pd.merge_asof(
@@ -420,7 +443,7 @@ def test_merge_asof_params_unsupported(left_right_timestamp_data):
         NotImplementedError,
         match=(
             "Snowpark pandas merge_asof method does not currently support parameters "
-            + "'left_index', 'right_index', 'suffixes', or 'tolerance'"
+            + "'suffixes', or 'tolerance'"
         ),
     ):
         pd.merge_asof(