SNOW-2090928: [Local Testing] Fix indexing issue in to_timestamp

Author: sfc-gh-jrose

src/snowflake/snowpark/mock/_functions.py

@@ -1065,10 +1065,13 @@ def convert_timestamp(row):
                 SnowparkLocalTestingException.raise_from_error(exc)
 
     res = column.to_frame().apply(convert_timestamp, axis=1).replace({pandas.NaT: None})
-    return [
-        x.to_pydatetime() if x is not None and hasattr(x, "to_pydatetime") else x
-        for x in res
-    ]
+    return pandas.Series(
+        [
+            x.to_pydatetime() if x is not None and hasattr(x, "to_pydatetime") else x
+            for x in res
+        ],
+        index=column.index,
+    )
 
 
 @patch("to_timestamp")
@@ -1102,6 +1105,7 @@ def mock_to_timestamp_ntz(
             TimestampType(TimestampTimeZone.NTZ), column.sf_type.nullable
         ),
         dtype=object,
+        index=result.index,
     )
 
 
@@ -1121,6 +1125,7 @@ def mock_to_timestamp_ltz(
             TimestampType(TimestampTimeZone.LTZ), column.sf_type.nullable
         ),
         dtype=object,
+        index=result.index,
     )
 
 
@@ -1132,12 +1137,14 @@ def mock_to_timestamp_tz(
 ):
     # _to_timestamp will use the tz present in the data.
     # Otherwise it adds an appropriate one by default.
+    result = _to_timestamp(column, fmt, try_cast, add_timezone=True)
     return ColumnEmulator(
-        data=_to_timestamp(column, fmt, try_cast, add_timezone=True),
+        data=result,
         sf_type=ColumnType(
             TimestampType(TimestampTimeZone.TZ), column.sf_type.nullable
         ),
         dtype=object,
+        index=result.index,
     )
 
 

tests/integ/scala/test_function_suite.py

@@ -84,6 +84,7 @@
     cume_dist,
     current_database,
     current_session,
+    date_format,
     date_part,
     date_trunc,
     dateadd,
@@ -1990,6 +1991,41 @@ def test_to_date(session):
     Utils.check_answer(df.select(to_date(df.a, "YYYY-MM-DD")), expected5)
 
 
+def test_date_format(session):
+    df = session.create_dataframe(
+        [
+            (date(2025, 1, 1), "A"),
+            (date(2025, 2, 2), "B"),
+            (date(2025, 3, 3), "A"),
+            (date(2025, 4, 4), "A"),
+        ],
+        ["date", "cat"],
+    )
+
+    # Check format works
+    Utils.check_answer(
+        df.withColumn("FORMATTED", date_format("DATE", "MMYYYY")),
+        [
+            Row(date(2025, 1, 1), "A", "012025"),
+            Row(date(2025, 2, 2), "B", "022025"),
+            Row(date(2025, 3, 3), "A", "032025"),
+            Row(date(2025, 4, 4), "A", "042025"),
+        ],
+    )
+
+    # Check format after filter works
+    Utils.check_answer(
+        df.filter(col("CAT") == "A").withColumn(
+            "FORMATTED", date_format("DATE", "MMYYYY")
+        ),
+        [
+            Row(date(2025, 1, 1), "A", "012025"),
+            Row(date(2025, 3, 3), "A", "032025"),
+            Row(date(2025, 4, 4), "A", "042025"),
+        ],
+    )
+
+
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="arrays_overlap is not yet supported in local testing mode.",