SNOW-2435494: Add support for iloc and head in faster pandas (#3910)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -122,6 +122,8 @@
   - `drop`
   - `invert`
   - `duplicated`
+  - `iloc`
+  - `head`
   - `columns` (e.g., df.columns = ["A", "B"])
   - `agg`
   - `min`

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

@@ -10762,6 +10762,36 @@ def take_2d_positional(
         self,
         index: Union["SnowflakeQueryCompiler", slice],
         columns: Union["SnowflakeQueryCompiler", slice, int, bool, list, AnyArrayLike],
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _take_2d_positional_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None and (
+            index == slice(None, None, None)
+            or (
+                isinstance(index, slice)
+                and (index.start is None or index.start == 0)
+                and (index.step is None or index.step == 1)
+            )
+        ):
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._take_2d_positional_internal(
+                    index=index,
+                    columns=columns,
+                )
+            )
+
+        qc = self._take_2d_positional_internal(
+            index=index,
+            columns=columns,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _take_2d_positional_internal(
+        self,
+        index: Union["SnowflakeQueryCompiler", slice],
+        columns: Union["SnowflakeQueryCompiler", slice, int, bool, list, AnyArrayLike],
     ) -> "SnowflakeQueryCompiler":
         """
         Index QueryCompiler with passed keys.

tests/integ/modin/test_faster_pandas.py

@@ -308,6 +308,49 @@ def test_duplicated(session):
     assert_series_equal(snow_result, native_result)
 
 
+@sql_count_checker(query_count=5)
+def test_iloc_head(session):
+    # create tables
+    table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+    session.create_dataframe(
+        native_pd.DataFrame([[1, 11], [2, 12], [3, 13]], columns=["A", "B"])
+    ).write.save_as_table(table_name, table_type="temp")
+
+    # create snow dataframes
+    df = pd.read_snowflake(table_name)
+    snow_result1 = df.iloc[:, [1]]
+    snow_result2 = df.iloc[0:2:1, [1]]
+    snow_result3 = df.head()
+
+    # verify that the input dataframe has a populated relaxed query compiler
+    assert df._query_compiler._relaxed_query_compiler is not None
+    assert df._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_result1._query_compiler._relaxed_query_compiler is not None
+    assert (
+        snow_result1._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    )
+    assert snow_result2._query_compiler._relaxed_query_compiler is not None
+    assert (
+        snow_result2._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    )
+    assert snow_result3._query_compiler._relaxed_query_compiler is not None
+    assert (
+        snow_result3._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    )
+
+    # create pandas dataframes
+    native_df = df.to_pandas()
+    native_result1 = native_df.iloc[:, [1]]
+    native_result2 = native_df.iloc[0:2:1, [1]]
+    native_result3 = native_df.head()
+
+    # compare results
+    assert_frame_equal(snow_result1, native_result1)
+    assert_frame_equal(snow_result2, native_result2)
+    assert_frame_equal(snow_result3, native_result3)
+
+
 @sql_count_checker(query_count=3)
 def test_invert(session):
     # create tables