diff --git a/src/snowflake/snowpark/_internal/utils.py b/src/snowflake/snowpark/_internal/utils.py index b61330783e..5578477ab6 100644 --- a/src/snowflake/snowpark/_internal/utils.py +++ b/src/snowflake/snowpark/_internal/utils.py @@ -63,7 +63,9 @@ from snowflake.connector._common.extras import ( MissingOptionalDependency, ModuleLikeObject, + pandas, pyarrow, + installed_pandas, installed_pyarrow, ) else: @@ -71,7 +73,9 @@ MissingOptionalDependency, MissingPandas, ModuleLikeObject, + pandas, pyarrow, + installed_pandas, ) # connector.options (v4) never exported installed_pyarrow as its own name. @@ -269,22 +273,6 @@ def _missing_pandas() -> MissingOptionalDependency: SUPPORTED_TABLE_TYPES = ["temp", "temporary", "transient"] -def _pandas_importer(): # noqa: E302 - """Helper function to lazily import pandas and return MissingPandas if not installed.""" - result = _missing_pandas() - try: - result = importlib.import_module("pandas") - # since we enable relative imports without dots this import gives us an issues when ran from test directory - from pandas import DataFrame # NOQA - except ImportError: # pragma: no cover - pass # pragma: no cover - return result - - -pandas = _pandas_importer() -installed_pandas = not isinstance(pandas, MissingOptionalDependency) - - class TempObjectType(Enum): TABLE = "TABLE" VIEW = "VIEW" diff --git a/src/snowflake/snowpark/mock/_options.py b/src/snowflake/snowpark/mock/_options.py index 7d1f088433..37d71e3d4f 100644 --- a/src/snowflake/snowpark/mock/_options.py +++ b/src/snowflake/snowpark/mock/_options.py @@ -4,7 +4,11 @@ import importlib -from snowflake.snowpark._internal.utils import MissingOptionalDependency, _missing_pandas +from snowflake.snowpark._internal.utils import ( + IS_V5_DRIVER, + MissingOptionalDependency, + _missing_pandas, +) try: import pandas @@ -15,15 +19,20 @@ installed_pandas = False -class MissingNumpy(MissingOptionalDependency): - """The class is specifically for numpy optional dependency.""" +if IS_V5_DRIVER: + from snowflake.connector._common.extras import numpy - _dep_name = "numpy" + installed_numpy = not isinstance(numpy, MissingOptionalDependency) +else: + class MissingNumpy(MissingOptionalDependency): + """The class is specifically for numpy optional dependency.""" -try: - numpy = importlib.import_module("numpy") - installed_numpy = True -except ImportError: - numpy = MissingNumpy() - installed_numpy = False + _dep_name = "numpy" + + try: + numpy = importlib.import_module("numpy") + installed_numpy = True + except ImportError: + numpy = MissingNumpy() + installed_numpy = False diff --git a/tests/unit/test_internal_utils.py b/tests/unit/test_internal_utils.py index f09cfc63d3..e3e6882dc5 100644 --- a/tests/unit/test_internal_utils.py +++ b/tests/unit/test_internal_utils.py @@ -4,11 +4,9 @@ import concurrent.futures import random import pytest -from snowflake.connector.options import MissingPandas from snowflake.snowpark._internal import utils from snowflake.snowpark._internal.utils import ( - _pandas_importer, generate_random_alphanumeric, split_snowflake_identifier_with_dot, ) @@ -148,16 +146,6 @@ def test_normalize_path_escapes_backslash_and_quote(raw_path, is_local): ), f"decoded={decoded!r} does not end with {expected_tail!r}" -def test__pandas_importer(): - imported_pandas = _pandas_importer() - try: - import pandas - - assert imported_pandas == pandas - except ImportError: - assert isinstance(imported_pandas, MissingPandas) - - def test_generate_random_alphanumeric(): random.seed(42) random_string1 = generate_random_alphanumeric()