SNOW-3217482: `mock_convert_timezone` fails on string inputs

[DPeled]

Please answer these questions before submitting your issue. Thanks!

1. What version of Python are you using?

   Python 3.11.13

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

    pandas==2.1.4
    pandas-gbq==0.31.0
    snowflake-snowpark-python==1.42.0
   

3. What did you do?

   When running convert_timezone in local testing with source_timestamp as string (varchar), the mock fails to run with the exception 'str' object has no attribute 'tzinfo'. It's because of that line.

   from snowflake.snowpark import Session
   from snowflake.snowpark.functions import lit, convert_timezone
   
   session = Session.builder.config('local_testing', True).create()
   df = session.create_dataframe([['2023-03-10T08:00:00+03:00']], schema=['date'])
   
   df.select(convert_timezone(lit('UTC'), 'date')).collect()

   The output is

   Traceback (most recent call last):
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_plan.py", line 661, in handle_function_expression
       result = func(*call_args[0], row_number=current_row, input_data=input_data)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_functions.py", line 125, in __call__
       result = self.impl(*args, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_functions.py", line 2049, in mock_convert_timezone
       res = combined.apply(_convert, axis=1)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/pandas/core/frame.py", line 10034, in apply
       return op.apply().__finalize__(self, method="apply")
              ^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/pandas/core/apply.py", line 837, in apply
       return self.apply_standard()
              ^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/pandas/core/apply.py", line 965, in apply_standard
       results, res_index = self.apply_series_generator()
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/pandas/core/apply.py", line 981, in apply_series_generator
       results[i] = self.func(v, *self.args, **self.kwargs)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_functions.py", line 2039, in _convert
       if source_time.tzinfo is None:
          ^^^^^^^^^^^^^^^^^^
   AttributeError: 'str' object has no attribute 'tzinfo'
   
   The above exception was the direct cause of the following exception:
   
   Traceback (most recent call last):
     File "<stdin>", line 1, in <module>
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/_internal/telemetry.py", line 295, in wrap
       result = func(*args, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/_internal/utils.py", line 1148, in call_wrapper
       return func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/dataframe.py", line 772, in collect
       return self._internal_collect_with_tag_no_telemetry(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/dataframe.py", line 845, in _internal_collect_with_tag_no_telemetry
       return self._session._conn.execute(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_connection.py", line 645, in execute
       res = execute_mock_plan(plan, plan.expr_to_alias)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_plan.py", line 1136, in execute_mock_plan
       column_series = calculate_expression(
                       ^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_plan.py", line 2261, in calculate_expression
       return calculate_expression(exp.child, input_data, analyzer, expr_to_alias)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_plan.py", line 2263, in calculate_expression
       return handle_function_expression(exp, input_data, analyzer, expr_to_alias)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/_plan.py", line 675, in handle_function_expression
       SnowparkLocalTestingException.raise_from_error(
     File "/Users/danpeled/Desktop/Canopy/data/airflow/dags/.venv/lib/python3.11/site-packages/snowflake/snowpark/mock/exceptions.py", line 20, in raise_from_error
       raise cls(
   snowflake.snowpark.mock.exceptions.SnowparkLocalTestingException: Error executing mocked function 'convert_timezone'. See error traceback for detailed information.

4. What did you expect to see?

   When executing in Snowflake

   SELECT convert_timezone('UTC', '2023-03-10T08:00:00+03:00')

   It returns 2023-03-10 05:00:00.000 +0000

It worth to mention that similar issue is mentioned in #1630

[sfc-gh-sghosh]

Hello @DPeled ,

Thanks for raising the issue.
We are able to reproduce the issue, this issue is different from previous one v.19.9 in #1630

The v1.19.0 fix for #1630 corrected the argument ordering but did not add string-to-datetime parsing which is the current issue.

As a workaround, could you please cast the string to a timestamp before calling convert_timezone

from snowflake.snowpark import Session
from snowflake.snowpark.functions import lit, convert_timezone,  col, to_timestamp

session = Session.builder.config('local_testing', True).create()
df = session.create_dataframe([['2023-03-10T08:00:00+03:00']], schema=['date'])

df.select(convert_timezone(lit('UTC'), to_timestamp(col('date')))).collect()
[Row(CONVERT_TIMEZONE('UTC', TO_TIMESTAMP("DATE"))=Timestamp('2023-03-10 02:30:00+0000', tz='tzfile('/usr/share/zoneinfo/UTC')'))]