Add a functionality in apply_in_pandas to support spark api (#3162)

Author: sfc-gh-dyadav

src/snowflake/snowpark/context.py

@@ -26,6 +26,9 @@
 _use_structured_type_semantics = False
 _use_structured_type_semantics_lock = threading.RLock()
 
+# This is an internal-only global flag, used to determine whether the api code which will be executed is compatible with snowflake.snowpark_connect
+_is_snowpark_connect_compatible_mode = False
+
 
 def _should_use_structured_type_semantics():
     global _use_structured_type_semantics

src/snowflake/snowpark/relational_grouped_dataframe.py

@@ -3,9 +3,12 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union
+import inspect
 
 import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto
+import snowflake.snowpark.context as context
 from snowflake.connector.options import pandas
+from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark import functions
 from snowflake.snowpark._internal.analyzer.expression import (
     Expression,
@@ -404,8 +407,36 @@ def apply_in_pandas(
             - :func:`~snowflake.snowpark.functions.pandas_udtf`
         """
 
+        partition_by = [Column(expr, _emit_ast=False) for expr in self._grouping_exprs]
+
+        # this is the case where this is being called from spark
+        # this is not handleing nested column access, it is assuming that the access in the function is not nested
+        original_columns: List[str] | None = None
+        key_columns: List[str] | None = None
+        if context._is_snowpark_connect_compatible_mode:
+            if self._dataframe._column_map is not None:
+                original_columns = [
+                    column.spark_name for column in self._dataframe._column_map.columns
+                ]
+            signature = inspect.signature(func)
+            parameters = signature.parameters
+            if len(parameters) == 2:
+                key_columns = [
+                    unquote_if_quoted(col.get_name()) for col in partition_by
+                ]
+
         class _ApplyInPandas:
             def end_partition(self, pdf: pandas.DataFrame) -> pandas.DataFrame:
+                if key_columns is not None:
+                    import numpy as np
+
+                    key_list = [pdf[key].iloc[0] for key in key_columns]
+                    numpy_array = np.array(key_list)
+                    keys = tuple(numpy_array)
+                if original_columns is not None:
+                    pdf.columns = original_columns
+                if key_columns is not None:
+                    return func(keys, pdf)
                 return func(pdf)
 
         # for vectorized UDTF
@@ -427,7 +458,6 @@ def end_partition(self, pdf: pandas.DataFrame) -> pandas.DataFrame:
             _emit_ast=_emit_ast,
             **kwargs,
         )
-        partition_by = [Column(expr, _emit_ast=False) for expr in self._grouping_exprs]
 
         df = self._dataframe.select(
             _apply_in_pandas_udtf(*self._dataframe.columns).over(

tests/integ/test_udtf.py

@@ -9,7 +9,7 @@
 
 import pytest
 
-from snowflake.snowpark import Row, Table
+from snowflake.snowpark import Row, Table, context
 from snowflake.snowpark._internal.utils import TempObjectType
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.functions import lit, udtf
@@ -532,6 +532,54 @@ def group_sum(pdf):
         ],
     )
 
+    class Column:
+        def __init__(self, spark_name: str) -> None:
+            self.spark_name = spark_name
+
+    class ColumnMap:
+        def __init__(self) -> None:
+            self.columns: List[Column] = []
+
+    # test with multiple columns in group by
+    df = session.createDataFrame(
+        [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v")
+    )
+
+    # this is to mock the current behavior
+    df._column_map = ColumnMap()
+    df._column_map.columns = [Column("id"), Column("v")]
+
+    context._is_snowpark_connect_compatible_mode = True
+
+    def normalize(pdf):
+        v = pdf.v
+        return pdf.assign(v=(v - v.mean()) / v.std())
+
+    df = (
+        df.group_by("id")
+        .applyInPandas(
+            normalize,
+            output_schema=StructType(
+                [
+                    StructField("id", IntegerType()),
+                    StructField("v", DoubleType()),
+                ]
+            ),
+        )
+        .orderBy(["id", "v"])
+    )
+
+    Utils.check_answer(
+        df,
+        [
+            Row(ID=1, V=-0.7071067811865475),
+            Row(ID=1, V=0.7071067811865475),
+            Row(ID=2, V=-0.8320502943378437),
+            Row(ID=2, V=-0.2773500981126146),
+            Row(ID=2, V=1.1094003924504583),
+        ],
+    )
+
 
 @pytest.mark.skipif(IS_IN_STORED_PROC, reason="Cannot create session in SP")
 def test_permanent_udtf_negative(session, db_parameters):