From 28cc861dcd2631d8b4dab79636f37c7f831fffbf Mon Sep 17 00:00:00 2001 From: Naren Krishna Date: Tue, 17 Sep 2024 08:35:22 -0700 Subject: [PATCH] SNOW-1662275: Add support for left_index and right_index for pd.merge_asof Signed-off-by: Naren Krishna --- CHANGELOG.md | 2 +- .../modin/supported/general_supported.rst | 3 +- .../compiler/snowflake_query_compiler.py | 4 +- tests/integ/modin/test_merge_asof.py | 43 ++++++++++++++----- 4 files changed, 37 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7048a3b728a..30715f5355f 100644 --- 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 index 95d9610202b..b3a71f023a9 100644 --- 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 index e971b15b6d6..a3981379aaf 100644 --- 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 index 51dda7889e7..5aab91fc9cb 100644 --- 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(