lint

Author: sfc-gh-jkew

src/snowflake/snowpark/modin/plugin/_internal/apply_utils.py

@@ -464,13 +464,13 @@ class ApplyFunc:
         def end_partition(self, df):  # type: ignore[no-untyped-def] # pragma: no cover
             # First column is row position, extract it for later use
             row_positions = df.iloc[:, 0]
-            
+
             # If we have index columns, set them as the index
             if num_index_columns > 0:
                 # Columns after row position are index columns, then data columns
-                index_cols = df.iloc[:, 1:1+num_index_columns]
-                data_cols = df.iloc[:, 1+num_index_columns:]
-                
+                index_cols = df.iloc[:, 1 : 1 + num_index_columns]
+                data_cols = df.iloc[:, 1 + num_index_columns :]
+
                 # Set the index using the index columns
                 if num_index_columns == 1:
                     index = index_cols.iloc[:, 0]
@@ -480,15 +480,17 @@ def end_partition(self, df):  # type: ignore[no-untyped-def] # pragma: no cover
                     # Multi-index case
                     index = native_pd.MultiIndex.from_arrays(
                         [index_cols.iloc[:, i] for i in range(num_index_columns)],
-                        names=index_column_pandas_labels if index_column_pandas_labels else None
+                        names=index_column_pandas_labels
+                        if index_column_pandas_labels
+                        else None,
                     )
                 data_cols.index = index
                 df = data_cols
             else:
                 # No index columns, use row position as index (original behavior)
                 df = df.iloc[:, 1:]
                 df.index = row_positions
-            
+
             df.columns = column_index
             df = df.apply(
                 func, axis=1, raw=raw, result_type=result_type, args=args, **kwargs
@@ -523,10 +525,10 @@ def end_partition(self, df):  # type: ignore[no-untyped-def] # pragma: no cover
             # - VALUE contains the result at this position.
             if isinstance(df, native_pd.DataFrame):
                 result = []
-                for idx, (row_position_index, series) in enumerate(df.iterrows()):
+                for idx, (_row_position_index, series) in enumerate(df.iterrows()):
                     # Use the actual row position from row_positions, not the index value
                     actual_row_position = row_positions.iloc[idx]
-                    
+
                     for i, (label, value) in enumerate(series.items()):
                         # If this is a tuple then we store each component with a 0-based
                         # lookup.  For example, (a,b,c) is stored as (0:a, 1:b, 2:c).

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

@@ -9602,7 +9602,9 @@ def _apply_with_udtf_and_dynamic_pivot_along_axis_1(
         # The apply function is encapsulated in a UDTF and run as a stored procedure on the pandas dataframe.
         # Determine if we should pass index columns to the UDTF
         # We pass index columns when the index is not the row position itself
-        index_columns_for_udtf = new_internal_df.index_column_snowflake_quoted_identifiers
+        index_columns_for_udtf = (
+            new_internal_df.index_column_snowflake_quoted_identifiers
+        )
         if row_position_snowflake_quoted_identifier in index_columns_for_udtf:
             # The row position IS the index (e.g., RangeIndex), don't pass index columns
             index_columns_for_udtf = []
@@ -9611,8 +9613,10 @@ def _apply_with_udtf_and_dynamic_pivot_along_axis_1(
         else:
             # Pass the actual index columns to the UDTF
             num_index_columns = len(index_columns_for_udtf)
-            index_column_pandas_labels_for_udtf = new_internal_df.index_column_pandas_labels
-        
+            index_column_pandas_labels_for_udtf = (
+                new_internal_df.index_column_pandas_labels
+            )
+
         func_udtf = create_udtf_for_apply_axis_1(
             row_position_snowflake_quoted_identifier,
             func,
@@ -10386,7 +10390,7 @@ def wrapped_func(*args, **kwargs):  # type: ignore[no-untyped-def] # pragma: no
             column_index = try_convert_index_to_native(
                 self._modin_frame.data_columns_index
             )
-            
+
             # get input types of index and data columns from the dataframe
             data_input_types = self._modin_frame.get_snowflake_type(
                 self._modin_frame.data_column_snowflake_quoted_identifiers

tests/integ/modin/frame/test_apply.py

@@ -147,38 +147,35 @@ def foo(row) -> str:
     with SqlCounter(query_count=4, join_count=0, udtf_count=0):
         eval_snowpark_pandas_result(snow_df, df, lambda x: x.apply(foo, axis=1))
 
-@pytest.mark.parametrize("index", [
-    None, 
-    ['a', 'b'], 
-    [100, 200]
-])
+
+@pytest.mark.parametrize("index", [None, ["a", "b"], [100, 200]])
 @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)
-    
+
     eval_snowpark_pandas_result(
-        snow_df, native_df, lambda x: x.apply(lambda row : row.name, axis=1)
+        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."""
     # Test with MultiIndex
-    multi_index = pd.MultiIndex.from_tuples([('A', 1), ('B', 2), ('C', 3)], names=['letter', 'number'])
+    multi_index = pd.MultiIndex.from_tuples(
+        [("A", 1), ("B", 2), ("C", 3)], names=["letter", "number"]
+    )
     native_df = native_pd.DataFrame([[1, 2], [3, 4], [5, 6]], index=multi_index)
     snow_df = pd.DataFrame(native_df)
-    
-    
+
     eval_snowpark_pandas_result(
-        snow_df, native_df, lambda x: x.apply(lambda row : row.name, axis=1)
+        snow_df, native_df, lambda x: x.apply(lambda row: row.name, axis=1)
     )
 
 
-
 @pytest.mark.xfail(strict=True, raises=NotImplementedError)
 @sql_count_checker(query_count=0)
 def test_frame_with_timedelta_index():