SNOW-2455993: Added support for concat in faster pandas (#3949)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -166,6 +166,7 @@
   - `groupby.nunique`
   - `groupby.size`
   - `drop_duplicates`
+  - `concat`
 - Reuse row count from the relaxed query compiler in `get_axis_len`.
 
 #### Bug Fixes

src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py

@@ -8614,6 +8614,57 @@ def concat(
         names: Optional[list[Hashable]] = None,
         verify_integrity: Optional[bool] = False,
         sort: Optional[bool] = False,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _concat_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if (
+            self._relaxed_query_compiler is not None
+            and all([qc._relaxed_query_compiler is not None for qc in other])
+            and axis == 0
+        ):
+            new_other = [
+                qc._relaxed_query_compiler
+                for qc in other
+                if qc._relaxed_query_compiler is not None
+            ]
+            relaxed_query_compiler = self._relaxed_query_compiler._concat_internal(
+                axis=axis,
+                other=new_other,
+                join=join,
+                ignore_index=ignore_index,
+                keys=keys,
+                levels=levels,
+                names=names,
+                verify_integrity=verify_integrity,
+                sort=sort,
+            )
+        qc = self._concat_internal(
+            axis=axis,
+            other=other,
+            join=join,
+            ignore_index=ignore_index,
+            keys=keys,
+            levels=levels,
+            names=names,
+            verify_integrity=verify_integrity,
+            sort=sort,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _concat_internal(
+        self,
+        axis: Axis,
+        other: list["SnowflakeQueryCompiler"],
+        *,
+        join: Optional[Literal["outer", "inner"]] = "outer",
+        ignore_index: bool = False,
+        keys: Optional[Sequence[Hashable]] = None,
+        levels: Optional[list[Sequence[Hashable]]] = None,
+        names: Optional[list[Hashable]] = None,
+        verify_integrity: Optional[bool] = False,
+        sort: Optional[bool] = False,
     ) -> "SnowflakeQueryCompiler":
         """
         Concatenate `self` with passed query compilers along specified axis.

tests/integ/modin/test_faster_pandas.py

@@ -250,6 +250,43 @@ def test_agg(session, func):
     assert_series_equal(snow_result4, native_result4, check_dtype=False)
 
 
+@sql_count_checker(query_count=5, union_count=1)
+def test_concat(session):
+    # create tables
+    table_name1 = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+    session.create_dataframe(
+        native_pd.DataFrame([[2, 12], [2, 12], [3, 13]], columns=["A", "B"])
+    ).write.save_as_table(table_name1, table_type="temp")
+    table_name2 = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+    session.create_dataframe(
+        native_pd.DataFrame([[2, 12], [2, 12], [3, 13]], columns=["A", "B"])
+    ).write.save_as_table(table_name2, table_type="temp")
+
+    # create snow dataframes
+    df1 = pd.read_snowflake(table_name1).sort_values("A", ignore_index=True)
+    df2 = pd.read_snowflake(table_name2).sort_values("A", ignore_index=True)
+    snow_result = pd.concat([df1, df2])
+
+    # verify that the input dataframe has a populated relaxed query compiler
+    assert df1._query_compiler._relaxed_query_compiler is not None
+    assert df1._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    assert df2._query_compiler._relaxed_query_compiler is not None
+    assert df2._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    # verify that the output dataframe also has a populated relaxed query compiler
+    assert snow_result._query_compiler._relaxed_query_compiler is not None
+    assert (
+        snow_result._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    )
+
+    # create pandas dataframes
+    native_df1 = df1.to_pandas()
+    native_df2 = df2.to_pandas()
+    native_result = native_pd.concat([native_df1, native_df2])
+
+    # compare results
+    assert_frame_equal(snow_result, native_result)
+
+
 @sql_count_checker(query_count=3)
 def test_drop(session):
     # create tables