Repository: modin-project/modin

Issue:
REFACTOR: Add merge_asof query compiler function



Discussion:

Resolving patch:
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -38,8 +38,9 @@
 from pandas.core.dtypes.common import is_scalar
 import pandas.core.resample
 import pandas
+from pandas._typing import IndexLabel, Suffixes
 import numpy as np
-from typing import List, Hashable
+from typing import List, Hashable, Optional
 
 
 def _get_axis(axis):
@@ -856,6 +857,143 @@ def merge(self, right, **kwargs):  # noqa: PR02
             self, right=right, **kwargs
         )
 
+    def _get_column_as_pandas_series(self, key):
+        """
+        Get column data by label as pandas.Series.
+
+        Parameters
+        ----------
+        key : Any
+            Column label.
+
+        Returns
+        -------
+        pandas.Series
+        """
+        result = self.getitem_array([key]).to_pandas().squeeze(axis=1)
+        if not isinstance(result, pandas.Series):
+            raise RuntimeError(
+                f"Expected getting column {key} to give "
+                + f"pandas.Series, but instead got {type(result)}"
+            )
+        return result
+
+    def merge_asof(
+        self,
+        right: "BaseQueryCompiler",
+        left_on: Optional[IndexLabel] = None,
+        right_on: Optional[IndexLabel] = None,
+        left_index: bool = False,
+        right_index: bool = False,
+        left_by=None,
+        right_by=None,
+        suffixes: Suffixes = ("_x", "_y"),
+        tolerance=None,
+        allow_exact_matches: bool = True,
+        direction: str = "backward",
+    ):
+        # Pandas fallbacks for tricky cases:
+        if (
+            # No idea how this works or why it does what it does; and in fact
+            # there's a Pandas bug suggesting it's wrong:
+            # https://github.com/pandas-dev/pandas/issues/33463
+            (left_index and right_on is not None)
+            # This is the case where by is a list of columns. If we're copying lots
+            # of columns out of Pandas, maybe not worth trying our path, it's not
+            # clear it's any better:
+            or not (left_by is None or is_scalar(left_by))
+            or not (right_by is None or is_scalar(right_by))
+            # The implementation below assumes that the right index is unique
+            # because it uses merge_asof to map each position in the merged
+            # index to the label of the one right row that should be merged
+            # at that row position.
+            or not right.index.is_unique
+        ):
+            return self.default_to_pandas(
+                pandas.merge_asof,
+                right,
+                left_on=left_on,
+                right_on=right_on,
+                left_index=left_index,
+                right_index=right_index,
+                left_by=left_by,
+                right_by=right_by,
+                suffixes=suffixes,
+                tolerance=tolerance,
+                allow_exact_matches=allow_exact_matches,
+                direction=direction,
+            )
+
+        if left_on is None:
+            left_column = self.index
+        else:
+            left_column = self._get_column_as_pandas_series(left_on)
+
+        if right_on is None:
+            right_column = right.index
+        else:
+            right_column = right._get_column_as_pandas_series(right_on)
+
+        left_pandas_limited = {"on": left_column}
+        right_pandas_limited = {"on": right_column, "right_labels": right.index}
+        extra_kwargs = {}  # extra arguments to Pandas merge_asof
+
+        if left_by is not None or right_by is not None:
+            extra_kwargs["by"] = "by"
+            left_pandas_limited["by"] = self._get_column_as_pandas_series(left_by)
+            right_pandas_limited["by"] = right._get_column_as_pandas_series(right_by)
+
+        # 1. Construct Pandas DataFrames with just the 'on' and optional 'by'
+        # columns, and the index as another column.
+        left_pandas_limited = pandas.DataFrame(left_pandas_limited, index=self.index)
+        right_pandas_limited = pandas.DataFrame(right_pandas_limited)
+
+        # 2. Use Pandas' merge_asof to figure out how to map labels on left to
+        # labels on the right.
+        merged = pandas.merge_asof(
+            left_pandas_limited,
+            right_pandas_limited,
+            on="on",
+            direction=direction,
+            allow_exact_matches=allow_exact_matches,
+            tolerance=tolerance,
+            **extra_kwargs,
+        )
+        # Now merged["right_labels"] shows which labels from right map to left's index.
+
+        # 3. Re-index right using the merged["right_labels"]; at this point right
+        # should be same length and (semantically) same order as left:
+        right_subset = right.reindex(
+            axis=0, labels=pandas.Index(merged["right_labels"])
+        )
+        if not right_index:
+            right_subset = right_subset.drop(columns=[right_on])
+        if right_by is not None and left_by == right_by:
+            right_subset = right_subset.drop(columns=[right_by])
+        right_subset.index = self.index
+
+        # 4. Merge left and the new shrunken right:
+        result = self.merge(
+            right_subset,
+            left_index=True,
+            right_index=True,
+            suffixes=suffixes,
+            how="left",
+        )
+
+        # 5. Clean up to match Pandas output:
+        if left_on is not None and right_index:
+            result = result.insert(
+                # In theory this could use get_indexer_for(), but that causes an error:
+                list(result.columns).index(left_on + suffixes[0]),
+                left_on,
+                result.getitem_array([left_on + suffixes[0]]),
+            )
+        if not left_index and not right_index:
+            result = result.reset_index(drop=True)
+
+        return result
+
     @doc_utils.add_refer_to("DataFrame.join")
     def join(self, right, **kwargs):  # noqa: PR02
         """
diff --git a/modin/pandas/general.py b/modin/pandas/general.py
--- a/modin/pandas/general.py
+++ b/modin/pandas/general.py
@@ -182,140 +182,38 @@ def merge_asof(
     ):
         raise ValueError("Can't combine left/right_index with left/right_on or on.")
 
-    # Pandas fallbacks for tricky cases:
-    if (
-        # No idea how this works or why it does what it does; and in fact
-        # there's a Pandas bug suggesting it's wrong:
-        # https://github.com/pandas-dev/pandas/issues/33463
-        (left_index and right_on is not None)
-        # This is the case where by is a list of columns. If we're copying lots
-        # of columns out of Pandas, maybe not worth trying our path, it's not
-        # clear it's any better:
-        or not isinstance(by, (str, type(None)))
-        or not isinstance(left_by, (str, type(None)))
-        or not isinstance(right_by, (str, type(None)))
-        # The implementation below assumes that the right index is unique
-        # because it uses merge_asof to map each position in the merged
-        # index to the label of the one right row that should be merged
-        # at that row position.
-        or not right.index.is_unique
-    ):
-        if isinstance(right, DataFrame):
-            right = to_pandas(right)
-        return DataFrame(
-            pandas.merge_asof(
-                to_pandas(left),
-                right,
-                on=on,
-                left_on=left_on,
-                right_on=right_on,
-                left_index=left_index,
-                right_index=right_index,
-                by=by,
-                left_by=left_by,
-                right_by=right_by,
-                suffixes=suffixes,
-                tolerance=tolerance,
-                allow_exact_matches=allow_exact_matches,
-                direction=direction,
-            )
-        )
-
-    left_column = None
-    right_column = None
-
     if on is not None:
         if left_on is not None or right_on is not None:
             raise ValueError("If 'on' is set, 'left_on' and 'right_on' can't be set.")
         left_on = on
         right_on = on
 
-    if left_on is not None:
-        left_column = to_pandas(left[left_on])
-    elif left_index:
-        left_column = left.index
-    else:
-        raise ValueError("Need some sort of 'on' spec")
-
-    if right_on is not None:
-        right_column = to_pandas(right[right_on])
-    elif right_index:
-        right_column = right.index
-    else:
-        raise ValueError("Need some sort of 'on' spec")
-
-    # If we haven't set these by now, there's a bug in this function.
-    assert left_column is not None
-    assert right_column is not None
-
     if by is not None:
         if left_by is not None or right_by is not None:
             raise ValueError("Can't have both 'by' and 'left_by' or 'right_by'")
         left_by = right_by = by
 
-    # List of columns case should have been handled by direct Pandas fallback
-    # earlier:
-    assert isinstance(left_by, (str, type(None)))
-    assert isinstance(right_by, (str, type(None)))
-
-    left_pandas_limited = {"on": left_column}
-    right_pandas_limited = {"on": right_column, "right_labels": right.index}
-    extra_kwargs = {}  # extra arguments to Pandas merge_asof
-
-    if left_by is not None or right_by is not None:
-        extra_kwargs["by"] = "by"
-        left_pandas_limited["by"] = to_pandas(left[left_by])
-        right_pandas_limited["by"] = to_pandas(right[right_by])
-
-    # 1. Construct Pandas DataFrames with just the 'on' and optional 'by'
-    # columns, and the index as another column.
-    left_pandas_limited = pandas.DataFrame(left_pandas_limited, index=left.index)
-    right_pandas_limited = pandas.DataFrame(right_pandas_limited)
-
-    # 2. Use Pandas' merge_asof to figure out how to map labels on left to
-    # labels on the right.
-    merged = pandas.merge_asof(
-        left_pandas_limited,
-        right_pandas_limited,
-        on="on",
-        direction=direction,
-        allow_exact_matches=allow_exact_matches,
-        tolerance=tolerance,
-        **extra_kwargs,
-    )
-    # Now merged["right_labels"] shows which labels from right map to left's index.
-
-    # 3. Re-index right using the merged["right_labels"]; at this point right
-    # should be same length and (semantically) same order as left:
-    right_subset = right.reindex(index=pandas.Index(merged["right_labels"]))
-    if not right_index:
-        right_subset.drop(columns=[right_on], inplace=True)
-    if right_by is not None and left_by == right_by:
-        right_subset.drop(columns=[right_by], inplace=True)
-    right_subset.index = left.index
-
-    # 4. Merge left and the new shrunken right:
-    result = merge(
-        left,
-        right_subset,
-        left_index=True,
-        right_index=True,
-        suffixes=suffixes,
-        how="left",
-    )
+    if left_on is None and not left_index:
+        raise ValueError("Must pass on, left_on, or left_index=True")
 
-    # 5. Clean up to match Pandas output:
-    if left_on is not None and right_index:
-        result.insert(
-            # In theory this could use get_indexer_for(), but that causes an error:
-            list(result.columns).index(left_on + suffixes[0]),
+    if right_on is None and not right_index:
+        raise ValueError("Must pass on, right_on, or right_index=True")
+
+    return DataFrame(
+        query_compiler=left._query_compiler.merge_asof(
+            right._query_compiler,
             left_on,
-            result[left_on + suffixes[0]],
+            right_on,
+            left_index,
+            right_index,
+            left_by,
+            right_by,
+            suffixes,
+            tolerance,
+            allow_exact_matches,
+            direction,
         )
-    if not left_index and not right_index:
-        result.index = pandas.RangeIndex(start=0, stop=len(result))
-
-    return result
+    )
 
 
 @_inherit_docstrings(pandas.pivot_table, apilink="pandas.pivot_table")


Test patch: