SNOW-2442822: Fix np.asarray(datetime_with_timezone_index) (#3963)

The current implementation of `Index.__array__` passes `self.dtype` to `self.to_pandas().__array__()`, but we should instead propagate the `dtype` argument from the `__array__` call.

Signed-off-by: sfc-gh-mvashishtha <[email protected]>

Author: sfc-gh-mvashishtha

CHANGELOG.md

@@ -41,6 +41,7 @@
 #### Bug Fixes
 
 - Fixed a bug in `DataFrameGroupBy.agg` where func is a list of tuples used to set the names of the output columns.
+- Fixed a bug where converting a modin datetime index with a timezone to a numpy array with `np.asarray` would cause a `TypeError`.
 
 #### Improvements
 

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

@@ -962,7 +962,7 @@ def get_valid_col_pos_list_from_columns(
 
         # convert float like keys to integers
         elif not is_integer_dtype(pos_array.dtype):
-            assert is_float_dtype(
+            assert pos_array.size == 0 or is_float_dtype(
                 pos_array.dtype
             ), "list-like key must be list of int or float"
             pos_list = pos_array.astype(int)

src/snowflake/snowpark/modin/plugin/extensions/index.py

@@ -917,10 +917,6 @@ def _summary(self, name: Any = None) -> str:
 
     @materialization_warning
     def __array__(self, dtype: Any = None) -> np.ndarray:
-        # Ensure that the existing index dtype is preserved in the returned array
-        # if no other dtype is given.
-        if dtype is None:
-            dtype = self.dtype
         return self.to_pandas().__array__(dtype=dtype)
 
     def __repr__(self) -> str:

tests/integ/modin/test_to_numpy.py

@@ -8,8 +8,10 @@
 import modin.pandas as pd
 import numpy as np
 import pandas as native_pd
+from pandas.api.types import is_datetime64tz_dtype
 import pytest
 from numpy.testing import assert_array_equal
+from pytest import param
 
 import snowflake.snowpark.modin.plugin  # noqa: F401
 from snowflake.snowpark._internal.utils import (
@@ -21,6 +23,10 @@
 from tests.utils import Utils
 
 
+def values_property_getter(object) -> np.ndarray:
+    return object.values
+
+
 @pytest.mark.parametrize(
     "data",
     [
@@ -40,11 +46,19 @@
         [datetime.time(1, 2, 3, 1), datetime.time(0, 0, 0), None],
         [datetime.datetime(2023, 1, 1), datetime.datetime(2023, 1, 1, 1, 2, 3)],
         [datetime.datetime(2023, 1, 1), datetime.datetime(2023, 1, 1, 1, 2, 3), None],
+        [pd.Timestamp(1, tz="UTC")],
     ],
 )
+@pytest.mark.parametrize(
+    "to_numpy",
+    (
+        param(lambda index: np.asarray(index), id="asarray"),
+        param(lambda index: index.to_numpy(), id="to_numpy"),
+        param(values_property_getter, id="values"),
+    ),
+)
 @pytest.mark.parametrize("pandas_obj", ["DataFrame", "Series", "Index"])
-@pytest.mark.parametrize("func", ["to_numpy", "values"])
-def test_to_numpy_basic(data, pandas_obj, func):
+def test_to_numpy_basic(data, pandas_obj, to_numpy, request):
     if pandas_obj == "Series":
         df = pd.Series(data)
         native_df = native_pd.Series(data)
@@ -54,16 +68,21 @@ def test_to_numpy_basic(data, pandas_obj, func):
     else:
         df = pd.DataFrame([data, data])
         native_df = native_pd.DataFrame([data, data])
-    with SqlCounter(query_count=1):
-        if func == "to_numpy":
-            assert_array_equal(df.to_numpy(), native_df.to_numpy())
-        else:
-            assert_array_equal(df.values, native_df.values)
+
+    with SqlCounter(
+        # modin_datetime_series_with_timezone.values internally calls .dtype,
+        # which triggers an extra groupby query to determine the dtype.
+        query_count=2
+        if to_numpy is values_property_getter
+        and pandas_obj == "Series"
+        and is_datetime64tz_dtype(native_df)
+        else 1
+    ):
+        snow_result = to_numpy(df)
+    native_result = to_numpy(native_df)
+    assert_array_equal(snow_result, native_result)
     if pandas_obj == "Series":
-        with SqlCounter(query_count=1):
-            res = df.to_numpy()
-        expected_res = native_df.to_numpy()
-        for r1, r2 in zip(res, expected_res):
+        for r1, r2 in zip(snow_result, native_result):
             # native pandas series returns a list of pandas Timestamp,
             # but Snowpark pandas returns a list of integers in ms.
             # Their values are equal