SNOW-2920252: [Local Testing] Lag window function produces ValueError when partitioning by multiple columns

[mykhailo-hurskyi]

1. What version of Python are you using?

   Python 3.9.6 (default, Oct 17 2025, 17:15:53)
   [Clang 17.0.0 (clang-1700.4.4.1)]

2. What are the Snowpark Python and pandas versions in the environment?

   pandas==2.3.3
   snowflake-snowpark-python==1.43.0

3. What did you do?

   When running the script below with the lag function, local-testing produces ValueError: 10 is not in list.
   This error only appears on a specific set of values as demonstrated in a script with "group_ok" and "group_error" columns.

   bug.py

4. What did you expect to see?

   When running the same snowpark code in Snowflake I get the following result:
   output.csv

   The code I used to run in Snowflake to verify expected behavior:
   snowflake_run.sql

Potential fix
There is a similar issue raised for other window functions like this one: #3459 which was followed by PR: #3462.

But it looks like lag and lead functions have a special handling logic in file _plan.py which was not touched in this PR. The issue looks very similar though and replacing res_index with pd_index in the line below fixed the issue locally for me.

[sfc-gh-sghosh]

Hello @mykhailo-hurskyi ,

Thanks for the issue.
We are able to reproduce and this is a existing known issue #3459
Will work on it to eliminate.

Regards,
Sujan