[SNOW-1638064] Support for asof join

src/snowflake/snowpark/dataframe.py

@@ -1631,6 +1631,7 @@ def sort(
             for c in _cols:
                 build_expr_from_snowpark_column_or_col_name(ast.cols.add(), c)
             ast.cols_variadic = is_variadic
+            self.set_ast_ref(ast.df)
 
         orders = []
         # `ascending` is represented by Expr in the AST.
@@ -3023,9 +3024,7 @@ def join(
                 elif isinstance(join_type, LeftAnti):
                     ast.join_type.sp_join_type__left_anti = True
                 elif isinstance(join_type, AsOf):
-                    raise NotImplementedError(
-                        "TODO SNOW-1638064: Add support for asof join to IR."
-                    )
+                    ast.join_type.sp_join_type__asof = True
                 else:
                     raise ValueError(f"Unsupported join type {join_type}")
 
@@ -4069,7 +4068,7 @@ def _show_string(self, n: int = 10, max_width: int = 50, **kwargs) -> str:
             # Phase 0 code where string gets formatted.
             if is_sql_select_statement(query):
                 result, meta = self._session._conn.get_result_and_metadata(
-                    self.limit(n)._plan, **kwargs
+                    self.limit(n, _emit_ast=False)._plan, **kwargs
                 )
             else:
                 res, meta = self._session._conn.get_result_and_metadata(

tests/ast/data/Dataframe.join.asof.test

@@ -0,0 +1,22 @@
+## TEST CASE
+
+df1 = session.create_dataframe([['A', 1, 15, 3.21],
+                               ['A', 2, 16, 3.22],
+                               ['B', 1, 17, 3.23],
+                               ['B', 2, 18, 4.23]],
+                              schema=["c1", "c2", "c3", "c4"])
+df2 = session.create_dataframe([['A', 1, 14, 3.19],
+                                ['B', 2, 16, 3.04]],
+                               schema=["c1", "c2", "c3", "c4"])
+
+df1.join(df2, on=(df1.c1 == df2.c1) & (df1.c2 == df2.c2), how="asof",
+         match_condition=(df1.c3 >= df2.c3), lsuffix="_L", rsuffix="_R") \
+         .order_by("C1_L", "C2_L").collect()
+
+## EXPECTED OUTPUT
+
+df1 = session.create_dataframe([["A", 1, 15, 3.21], ["A", 2, 16, 3.22], ["B", 1, 17, 3.23], ["B", 2, 18, 4.23]], schema=["c1", "c2", "c3", "c4"])
+
+df2 = session.create_dataframe([["A", 1, 14, 3.19], ["B", 2, 16, 3.04]], schema=["c1", "c2", "c3", "c4"])
+
+df1.join(df2, on=(df1["c1"] == df2["c1"]) & (df1["c2"] == df2["c2"]), how="asof").sort("C1_L", "C2_L").collect()