Revert back to other verified unnecessary mock modules

Author: sfc-gh-nqi

src/snowflake/snowpark/mock/_connection.py

@@ -11,18 +11,7 @@
 from copy import copy
 from decimal import Decimal
 from logging import getLogger
-from typing import (
-    IO,
-    Any,
-    Dict,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    Tuple,
-    Union,
-    TYPE_CHECKING,
-)
+from typing import IO, Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union
 from unittest.mock import Mock
 
 import snowflake.snowpark.mock._constants
@@ -47,6 +36,7 @@
 )
 from snowflake.snowpark.async_job import AsyncJob, _AsyncResultType
 from snowflake.snowpark.exceptions import SnowparkSessionException
+from snowflake.snowpark.mock._options import pandas
 from snowflake.snowpark.mock._plan import MockExecutionPlan, execute_mock_plan
 from snowflake.snowpark.mock._snowflake_data_type import ColumnEmulator, TableEmulator
 from snowflake.snowpark.mock._stage_registry import StageEntityRegistry
@@ -69,9 +59,6 @@
 PARAM_INTERNAL_APPLICATION_NAME = "internal_application_name"
 PARAM_INTERNAL_APPLICATION_VERSION = "internal_application_version"
 
-if TYPE_CHECKING:
-    from snowflake.snowpark.mock._options import pandas
-
 
 class MockedSnowflakeConnection(SnowflakeConnection):
     def __init__(self, *args, **kwargs) -> None:
@@ -157,8 +144,6 @@ def write_table(
             mode: SaveMode,
             column_names: Optional[List[str]] = None,
         ) -> List[Row]:
-            from snowflake.snowpark.mock._options import pandas
-
             with self._lock:
                 for column in table.columns:
                     if (
@@ -719,8 +704,6 @@ def execute(
             rows = [r for r in res]
 
         if to_pandas:
-            from snowflake.snowpark.mock._options import pandas
-
             pandas_df = pandas.DataFrame()
             for col_name in res.columns:
                 pandas_df[unquote_if_quoted(col_name)] = res[col_name].tolist()
@@ -826,8 +809,6 @@ def max_string_size(self) -> int:
 
 
 def _fix_pandas_df_fixed_type(table_res: TableEmulator) -> "pandas.DataFrame":
-    from snowflake.snowpark.mock._options import pandas
-
     pd_df = pandas.DataFrame()
     for col_name in table_res.columns:
         col_sf_type = table_res.sf_types[col_name]

src/snowflake/snowpark/mock/_functions.py

@@ -24,7 +24,7 @@
 import snowflake.snowpark
 from snowflake.snowpark._internal.analyzer.expression import FunctionExpression
 from snowflake.snowpark._internal.utils import unalias_datetime_part
-from snowflake.snowpark.mock._options import numpy
+from snowflake.snowpark.mock._options import numpy, pandas
 from snowflake.snowpark.mock._snowflake_data_type import (
     _TIMESTAMP_TYPE_MAPPING,
     _TIMESTAMP_TYPE_TIMEZONE_MAPPING,
@@ -500,8 +500,6 @@ def mock_array_agg(column: ColumnEmulator, is_distinct: bool) -> ColumnEmulator:
 
 @patch("array_construct")
 def mock_array_construct(*columns):
-    from snowflake.snowpark.mock._options import pandas
-
     if len(columns) == 0:
         data = [[]]
     else:
@@ -1066,8 +1064,6 @@ def convert_timestamp(row):
             else:
                 SnowparkLocalTestingException.raise_from_error(exc)
 
-    from snowflake.snowpark.mock._options import pandas
-
     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
@@ -1454,8 +1450,6 @@ def mock_to_binary(
 
 @patch("coalesce")
 def mock_coalesce(*exprs):
-    from snowflake.snowpark.mock._options import pandas
-
     if len(exprs) < 2:
         raise SnowparkLocalTestingException(
             f"not enough arguments for function [COALESCE], got {len(exprs)}, expected at least two"
@@ -1622,8 +1616,6 @@ def construct_dict(x):
             if x[i] is not None and not (drop_nulls and x[i + 1] is None)
         }
 
-    from snowflake.snowpark.mock._options import pandas
-
     combined = pandas.concat(exprs, axis=1, ignore_index=True)
     return combined.apply(construct_dict, axis=1)
 
@@ -1653,8 +1645,6 @@ def add_years(date, duration):
 
 
 def add_months(scalar, date, duration):
-    from snowflake.snowpark.mock._options import pandas
-
     res = (
         pandas.to_datetime(date) + pandas.DateOffset(months=scalar * duration)
     ).to_pydatetime()
@@ -1723,7 +1713,6 @@ def mock_date_part(part: str, datetime_expr: ColumnEmulator):
     """
     unaliased = unalias_datetime_part(part)
     datatype = datetime_expr.sf_type.datatype
-    from snowflake.snowpark.mock._options import pandas
 
     # Year of week is another alias unique to date_part
     if unaliased == "yearofweek":
@@ -1803,8 +1792,6 @@ def mock_date_trunc(part: str, datetime_expr: ColumnEmulator) -> ColumnEmulator:
     """
     # Map snowflake time unit to pandas rounding alias
     # Not all units have an alias so handle those with a special case
-    from snowflake.snowpark.mock._options import pandas
-
     SUPPORTED_UNITS = {
         "day": "D",
         "hour": "h",
@@ -1903,7 +1890,6 @@ def func(x, y):
     data = []
     for x, y in zip(col1, col2):
         data.append(None if x is None or y is None else func(x, y))
-    from snowflake.snowpark.mock._options import pandas
 
     return ColumnEmulator(
         pandas.Series(data, dtype=object),
@@ -2015,7 +2001,6 @@ def mock_convert_timezone(
     For timezone information, refer to the `Snowflake SQL convert_timezone notes <https://docs.snowflake.com/en/sql-reference/functions/convert_timezone.html#usage-notes>`_
     """
     import dateutil
-    from snowflake.snowpark.mock._options import pandas
 
     # mock_convert_timezone matches the sql function call semantics.
     # It has different parameters when called with 2 or 3 args.
@@ -2117,7 +2102,6 @@ def mock_concat(*columns: ColumnEmulator) -> ColumnEmulator:
         SnowparkLocalTestingException.raise_from_error(
             ValueError("concat expects one or more column(s) to be passed in.")
         )
-    from snowflake.snowpark.mock._options import pandas
 
     pdf = pandas.concat(columns, axis=1).reset_index(drop=True)
     result = pdf.T.apply(
@@ -2135,7 +2119,6 @@ def mock_concat_ws(*columns: ColumnEmulator) -> ColumnEmulator:
                 "concat_ws expects a seperator column and one or more value column(s) to be passed in."
             )
         )
-    from snowflake.snowpark.mock._options import pandas
 
     # Don't reset index, to preserve original indices from filtered DataFrames
     pdf = pandas.concat(columns, axis=1)
@@ -2258,7 +2241,6 @@ def _rank(raw_input, dense=False):
                 rank = index
         previous = value
         final_values.append(rank)
-    from snowflake.snowpark.mock._options import pandas
 
     return pandas.Series(final_values, index=raw_input.index)
 

src/snowflake/snowpark/mock/_nop_connection.py

@@ -16,6 +16,8 @@
     Union,
 )
 
+from snowflake.snowpark.mock._options import pandas
+
 from snowflake.connector.cursor import ResultMetadata
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark._internal.analyzer.expression import Attribute
@@ -57,8 +59,6 @@
 )
 
 if TYPE_CHECKING:
-    from snowflake.snowpark.mock._options import pandas
-
     try:
         from snowflake.connector.cursor import ResultMetadataV2
     except ImportError:
@@ -123,7 +123,6 @@ def execute(
         List[Row], "pandas.DataFrame", Iterator[Row], Iterator["pandas.DataFrame"]
     ]:
         source_plan = plan.source_plan
-        from snowflake.snowpark.mock._options import pandas
 
         if hasattr(source_plan, "execution_queries"):
             # If temp read-only table, explicitly create it.

src/snowflake/snowpark/mock/_pandas_util.py

@@ -9,6 +9,7 @@
     quote_name_without_upper_casing,
 )
 from snowflake.snowpark._internal.type_utils import infer_type
+from snowflake.snowpark.mock._options import pandas as pd
 from snowflake.snowpark.mock._telemetry import LocalTestOOBTelemetryService
 from snowflake.snowpark.table import Table
 from snowflake.snowpark.types import (
@@ -28,7 +29,6 @@
 )
 
 if TYPE_CHECKING:
-    from snowflake.snowpark.mock._options import pandas as pd
     from snowflake.snowpark import DataFrame, Session
 
 
@@ -43,7 +43,6 @@ def _extract_schema_and_data_from_pandas_df(
     pandas type related doc: https://pandas.pydata.org/docs/user_guide/basics.html#dtypes
     """
     import numpy
-    from snowflake.snowpark.mock._options import pandas as pd
 
     # Ensure pandas types are fully loaded before checking them
     def _ensure_pandas_types_available():

src/snowflake/snowpark/mock/_plan.py

@@ -42,6 +42,12 @@
 )
 from snowflake.snowpark.mock._udf_utils import coerce_variant_input, remove_null_wrapper
 from snowflake.snowpark.mock._util import ImportContext, get_fully_qualified_name
+from snowflake.snowpark.mock._window_utils import (
+    EntireWindowIndexer,
+    RowFrameIndexer,
+    is_rank_related_window_function,
+)
+
 from snowflake.snowpark.mock.exceptions import SnowparkLocalTestingException
 
 if TYPE_CHECKING:
@@ -142,6 +148,7 @@
 )
 from snowflake.snowpark.column import Column
 from snowflake.snowpark.mock._functions import MockedFunctionRegistry, cast_column_to
+from snowflake.snowpark.mock._options import pandas as pd
 from snowflake.snowpark.mock._select_statement import (
     MockSelectable,
     MockSelectableEntity,
@@ -178,9 +185,6 @@
     _NumericType,
 )
 
-if TYPE_CHECKING:
-    from snowflake.snowpark.mock._options import pandas as pd
-
 
 class MockExecutionPlan(LogicalPlan):
     def __init__(
@@ -427,12 +431,6 @@ def search_boundary_idx(idx, delta, _win):
         return idx
 
     if order_spec:
-        from snowflake.snowpark.mock._options import pandas as pd
-
-        # Use factory function to get proper indexer class
-        from snowflake.snowpark.mock._window_utils import _get_entire_window_indexer
-
-        EntireWindowIndexer = _get_entire_window_indexer()
         ordered_windows = [
             handle_order_by_clause(order_spec, win, analyzer, expr_to_alias)
             for win in res.rolling(EntireWindowIndexer())
@@ -518,10 +516,6 @@ def search_boundary_idx(idx, delta, _win):
 
                 windows.append(win[start_idx : end_idx + 1])
     else:  # If order by is not specified, just use the entire window
-        # Use factory function to get proper indexer class
-        from snowflake.snowpark.mock._window_utils import _get_entire_window_indexer
-
-        EntireWindowIndexer = _get_entire_window_indexer()
         windows = res.rolling(EntireWindowIndexer())
     return windows
 
@@ -838,7 +832,6 @@ def handle_udtf_expression(
 ):
 
     # TODO: handle and support imports + other udtf attributes.
-    from snowflake.snowpark.mock._options import pandas as pd
 
     udtf_registry = analyzer.session.udtf
     udtf_name = exp.func_name
@@ -1031,7 +1024,6 @@ def execute_mock_plan(
     expr_to_alias: Optional[Dict[str, str]] = None,
 ) -> Union[TableEmulator, List[Row]]:
     import numpy as np
-    from snowflake.snowpark.mock._options import pandas as pd
 
     if expr_to_alias is None:
         expr_to_alias = plan.expr_to_alias
@@ -2212,7 +2204,6 @@ def calculate_expression(
     setting keep_literal to false returns a ColumnEmulator wrapping the Python datatype of a Literal
     """
     import numpy as np
-    from snowflake.snowpark.mock._options import pandas as pd
 
     registry = MockedFunctionRegistry.get_or_create()
 
@@ -2571,13 +2562,6 @@ def _match_pattern(row) -> bool:
 
         return output_data
     if isinstance(exp, WindowExpression):
-        # Lazy import to avoid importing pandas at module level
-        from snowflake.snowpark.mock._window_utils import (
-            _get_entire_window_indexer,
-            _get_row_frame_indexer,
-            is_rank_related_window_function,
-        )
-
         window_function = exp.window_function
         window_spec = exp.window_spec
 
@@ -2642,8 +2626,6 @@ def _match_pattern(row) -> bool:
                 if not isinstance(windows, list):
                     pd_index = list(windows.count().index)
             else:
-                # Create indexer class dynamically to avoid module-level pandas import
-                EntireWindowIndexer = _get_entire_window_indexer()
                 indexer = EntireWindowIndexer()
                 rolling = res.rolling(indexer)
                 windows = [ordered.loc[w.index] for w in rolling]
@@ -2652,8 +2634,6 @@ def _match_pattern(row) -> bool:
                 pd_index = list(rolling.count().index)
 
         elif isinstance(window_spec.frame_spec.frame_type, RowFrame):
-            # Create indexer class dynamically to avoid module-level pandas import
-            RowFrameIndexer = _get_row_frame_indexer()
             indexer = RowFrameIndexer(frame_spec=window_spec.frame_spec)
             res = res.rolling(indexer)
             res_index = list(res.count().index)

src/snowflake/snowpark/mock/_stage_registry.py

@@ -22,6 +22,7 @@
     unwrap_stage_location_single_quote,
 )
 from snowflake.snowpark.mock._functions import mock_to_char
+from snowflake.snowpark.mock._options import pandas as pd
 from snowflake.snowpark.mock._snowflake_data_type import (
     ColumnEmulator,
     ColumnType,
@@ -409,7 +410,6 @@ def read_file(
         options: Dict[str, str],
     ) -> TableEmulator:
         from snowflake.snowpark.mock import CUSTOM_JSON_DECODER
-        from snowflake.snowpark.mock._options import pandas as pd
 
         stage_source_dir_path = os.path.join(self._working_directory, stage_location)
 

src/snowflake/snowpark/mock/_util.py

@@ -8,9 +8,10 @@
 import sys
 import tempfile
 from functools import cmp_to_key, partial
-from typing import Any, Callable, Iterable, Optional, Set, Tuple, Union, TYPE_CHECKING
+from typing import Any, Callable, Iterable, Optional, Set, Tuple, Union
 
 from snowflake.snowpark._internal.utils import parse_table_name, quote_name
+from snowflake.snowpark.mock._options import pandas as pd
 from snowflake.snowpark.mock._snowflake_data_type import ColumnEmulator
 from snowflake.snowpark.types import (
     ArrayType,
@@ -33,9 +34,6 @@
     _NumericType,
 )
 
-if TYPE_CHECKING:
-    from snowflake.snowpark.mock._options import pandas as pd
-
 
 # placeholder map helps convert wildcard to reg. In practice, we convert wildcard to a middle string first,
 # and then convert middle string to regex. See the following example: