Add test for an index from column

Author: sfc-gh-jkew

tests/integ/modin/frame/test_apply.py

@@ -152,7 +152,6 @@ def foo(row) -> str:
 @sql_count_checker(query_count=5, join_count=2, udtf_count=1)
 def test_apply_axis_1_index_preservation(index):
     """Test that apply(axis=1) preserves index values correctly."""
-    # Test with default RangeIndex
     native_df = native_pd.DataFrame([[1, 2], [3, 4]], index=index)
     snow_df = pd.DataFrame(native_df)
 
@@ -161,6 +160,21 @@ def test_apply_axis_1_index_preservation(index):
     )
 
 
+@sql_count_checker(query_count=5, join_count=2, udtf_count=1)
+def test_apply_axis_1_index_from_col():
+    """Test that apply(axis=1) preserves an index when set from a column"""
+    native_df = native_pd.DataFrame(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]], columns=["a", "b", "c"]
+    )
+    snow_df = pd.DataFrame(native_df)
+    snow_df = snow_df.set_index("a")
+    native_df = native_df.set_index("a")
+
+    eval_snowpark_pandas_result(
+        snow_df, native_df, lambda x: x.apply(lambda row: row.name, axis=1)
+    )
+
+
 @sql_count_checker(query_count=5, join_count=2, udtf_count=1)
 def test_apply_axis_1_multiindex_preservation():
     """Test that apply(axis=1) preserves MultiIndex values correctly."""