diff --git a/DESCRIPTION.md b/DESCRIPTION.md index 6b2be621f5..f78b513b7f 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -8,6 +8,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne # Release Notes - NEXT_RELEASE(TBD) + - Added a new `arrow` extra (`pip install "snowflake-connector-python[arrow]"`) that installs only PyArrow, allowing `fetch_arrow_all` and `fetch_arrow_batches` to be used without pandas installed. The `pandas` extra now depends on the `arrow` extra and installs the same packages as before (SNOW-710684). - Fixed `split_statements` treating `//` as SQL instead of a line comment, which could merge multiple statements when a `//` comment contained an apostrophe (SNOW-3772985). - Fixed large-file PUT uploads to internal Azure stages failing against the Azure 50,000-block-per-blob limit. The Azure multipart chunk size is now scaled up dynamically for very large files (mirroring the existing S3 behavior), and the default Azure chunk size was raised from 4 MB to 8 MB (consistent with S3) for better throughput (SNOW-3839943). diff --git a/setup.cfg b/setup.cfg index 896f869a36..598bcf3e4d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -86,6 +86,9 @@ console_scripts = snowflake-dump-certs = snowflake.connector.tool.dump_certs:main [options.extras_require] +arrow = + pyarrow>=14.0.1,<24; python_version >= '3.14' + pyarrow>=14.0.1; python_version < '3.14' boto = boto3>=1.24 botocore>=1.24 @@ -111,7 +114,6 @@ development = pandas = pandas>=1.0.0,<3.0.0; python_version < '3.13' pandas>=2.1.2,<3.0.0; python_version >= '3.13' - pyarrow>=14.0.1,<24; python_version >= '3.14' - pyarrow>=14.0.1; python_version < '3.14' + snowflake-connector-python[arrow] secure-local-storage = keyring>=23.1.0,<26.0.0 diff --git a/src/snowflake/connector/options.py b/src/snowflake/connector/options.py index e859ac68dd..24cf7de141 100644 --- a/src/snowflake/connector/options.py +++ b/src/snowflake/connector/options.py @@ -42,6 +42,12 @@ class MissingPandas(MissingOptionalDependency): _dep_name = "pandas" +class MissingPyArrow(MissingOptionalDependency): + """The class is specifically for pyarrow optional dependency.""" + + _dep_name = "pyarrow" + + class MissingKeyring(MissingOptionalDependency): """The class is specifically for sso optional dependency.""" @@ -91,19 +97,33 @@ def warn_incompatible_dep( ) -def _import_or_missing_pandas_option() -> ( - tuple[ModuleLikeObject, ModuleLikeObject, bool] -): - """This function tries importing the following packages: pandas, pyarrow. +def _import_or_missing_pandas_option( + installed_pyarrow: bool, +) -> tuple[ModuleLikeObject, bool]: + """This function tries importing pandas. - If available it returns pandas and pyarrow packages with a flag of whether they were imported. - It also warns users if they have an unsupported pyarrow version installed if possible. + Pandas fetch APIs also require pyarrow, so pandas is only reported as + available when pyarrow is too. """ + if not installed_pyarrow: + return MissingPandas(), False try: pandas = 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 + return pandas, True + except ImportError: + return MissingPandas(), False + + +def _import_or_missing_pyarrow_option() -> tuple[ModuleLikeObject, bool]: + """This function tries importing pyarrow. + + If available it returns the pyarrow package with a flag of whether it was imported. + It also warns users if they have an unsupported pyarrow version installed if possible. + """ + try: pyarrow = importlib.import_module("pyarrow") # set default memory pool to system for pyarrow to_pandas conversion @@ -118,30 +138,33 @@ def _import_or_missing_pandas_option() -> ( dependencies = snowflake_connector_dist.metadata.get_all( "Requires-Dist", [] ) - pandas_pyarrow_extra = None + arrow_pyarrow_extra = None for dependency in dependencies: dep = Requirement(dependency) if ( dep.marker is not None - and dep.marker.evaluate({"extra": "pandas"}) + and dep.marker.evaluate({"extra": "arrow"}) and dep.name == "pyarrow" ): - pandas_pyarrow_extra = dep + arrow_pyarrow_extra = dep break installed_pyarrow_version = pyarrow_dist.version - if not pandas_pyarrow_extra.specifier.contains(installed_pyarrow_version): + # metadata from installs predating the arrow extra has no such pin + if arrow_pyarrow_extra is not None and not ( + arrow_pyarrow_extra.specifier.contains(installed_pyarrow_version) + ): warn_incompatible_dep( - "pyarrow", installed_pyarrow_version, pandas_pyarrow_extra + "pyarrow", installed_pyarrow_version, arrow_pyarrow_extra ) except PackageNotFoundError as e: logger.info( f"Cannot determine if compatible pyarrow is installed because of missing package(s): {e}" ) - return pandas, pyarrow, True + return pyarrow, True except ImportError: - return MissingPandas(), MissingPandas(), False + return MissingPyArrow(), False def _import_or_missing_keyring_option() -> tuple[ModuleLikeObject, bool]: @@ -191,7 +214,8 @@ def _import_or_missing_azure_identity_option() -> ( # Create actual constants to be imported from this file -pandas, pyarrow, installed_pandas = _import_or_missing_pandas_option() +pyarrow, installed_pyarrow = _import_or_missing_pyarrow_option() +pandas, installed_pandas = _import_or_missing_pandas_option(installed_pyarrow) keyring, installed_keyring = _import_or_missing_keyring_option() botocore, boto3, installed_boto = _import_or_missing_boto_option() aiobotocore, aioboto3, installed_aioboto = _import_or_missing_aioboto_option() diff --git a/src/snowflake/connector/result_batch.py b/src/snowflake/connector/result_batch.py index f54beb8acb..d3cd5bb178 100644 --- a/src/snowflake/connector/result_batch.py +++ b/src/snowflake/connector/result_batch.py @@ -7,7 +7,7 @@ from base64 import b64decode from enum import Enum, unique from logging import getLogger -from typing import TYPE_CHECKING, Any, Callable, Iterator, NamedTuple, Sequence +from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, Sequence from typing_extensions import Self @@ -37,7 +37,7 @@ if TYPE_CHECKING: # pragma: no cover from pandas import DataFrame - from pyarrow import DataType, Table + from pyarrow import Table from .connection import SnowflakeConnection from .converter import SnowflakeConverterType @@ -45,9 +45,6 @@ from .vendored.requests import Response -# emtpy pyarrow type array corresponding to FIELD_TYPES -FIELD_TYPE_TO_PA_TYPE: list[Callable[[ResultMetadataV2], DataType]] = [] - # qrmk related constants SSE_C_ALGORITHM = "x-amz-server-side-encryption-customer-algorithm" SSE_C_KEY = "x-amz-server-side-encryption-customer-key" @@ -824,12 +821,8 @@ def _get_arrow_iter( def _create_empty_table(self) -> Table: """Returns empty Arrow table based on schema""" - if installed_pandas: - # initialize pyarrow type array corresponding to FIELD_TYPES - FIELD_TYPE_TO_PA_TYPE = [e.pa_type for e in FIELD_TYPES] fields = [ - pa.field(s.name, FIELD_TYPE_TO_PA_TYPE[s.type_code](s)) - for s in self._schema + pa.field(s.name, FIELD_TYPES[s.type_code].pa_type(s)) for s in self._schema ] return pa.schema(fields).empty_table() diff --git a/test/integ/pandas_it/test_unit_options.py b/test/integ/pandas_it/test_unit_options.py index 9038e98d7c..bb30c6885e 100644 --- a/test/integ/pandas_it/test_unit_options.py +++ b/test/integ/pandas_it/test_unit_options.py @@ -1,27 +1,28 @@ from __future__ import annotations import logging +import warnings from unittest import mock import pytest try: from snowflake.connector.options import ( - MissingPandas, - _import_or_missing_pandas_option, + MissingPyArrow, + _import_or_missing_pyarrow_option, ) except ImportError: - MissingPandas = None - _import_or_missing_pandas_option = None + MissingPyArrow = None + _import_or_missing_pyarrow_option = None from importlib.metadata import PackageNotFoundError, distribution @pytest.mark.skipif( - MissingPandas is None or _import_or_missing_pandas_option is None, + MissingPyArrow is None or _import_or_missing_pyarrow_option is None, reason="No snowflake.connector.options is available. It can be the case if running old driver tests", ) -def test_pandas_option_reporting(caplog): +def test_pyarrow_option_reporting(caplog): """Tests for the weird case where someone can import pyarrow, but setuptools doesn't know about it. This issue was brought to attention in: https://github.com/snowflakedb/snowflake-connector-python/issues/412 @@ -37,12 +38,68 @@ def modified_distribution(name, *args, **kwargs): wraps=modified_distribution, ): caplog.set_level(logging.DEBUG, "snowflake.connector") - pandas, pyarrow, installed_pandas = _import_or_missing_pandas_option() - assert installed_pandas - assert not isinstance(pandas, MissingPandas) - assert not isinstance(pyarrow, MissingPandas) + pyarrow, installed_pyarrow = _import_or_missing_pyarrow_option() + assert installed_pyarrow + assert not isinstance(pyarrow, MissingPyArrow) assert ( "Cannot determine if compatible pyarrow is installed because of missing package(s)" in caplog.text ) assert "TestErrorMessage" in caplog.text + + +@pytest.mark.skipif( + MissingPyArrow is None or _import_or_missing_pyarrow_option is None, + reason="No snowflake.connector.options is available. It can be the case if running old driver tests", +) +def test_pyarrow_version_check_reads_arrow_extra(): + """The supported pyarrow range comes from the arrow extra's pin.""" + pyarrow_dist = mock.Mock(version="2.0.0") + connector_dist = mock.Mock() + connector_dist.metadata.get_all.return_value = [ + 'pyarrow>=1; extra == "pandas"', + 'pyarrow<1; extra == "arrow"', + ] + + def modified_distribution(name): + if name == "pyarrow": + return pyarrow_dist + if name == "snowflake-connector-python": + return connector_dist + return distribution(name) + + with mock.patch( + "snowflake.connector.options.distribution", + wraps=modified_distribution, + ): + with pytest.warns(UserWarning, match="pyarrow<1"): + pyarrow, installed_pyarrow = _import_or_missing_pyarrow_option() + + assert installed_pyarrow + assert not isinstance(pyarrow, MissingPyArrow) + + +@pytest.mark.skipif( + MissingPyArrow is None or _import_or_missing_pyarrow_option is None, + reason="No snowflake.connector.options is available. It can be the case if running old driver tests", +) +def test_pyarrow_version_check_tolerates_metadata_without_arrow_extra(): + """Metadata predating the arrow extra must not warn or break the import.""" + connector_dist = mock.Mock() + connector_dist.metadata.get_all.return_value = ['pyarrow>=1; extra == "pandas"'] + + def modified_distribution(name): + if name == "snowflake-connector-python": + return connector_dist + return distribution(name) + + with mock.patch( + "snowflake.connector.options.distribution", + wraps=modified_distribution, + ): + with warnings.catch_warnings(): + warnings.simplefilter("error") + pyarrow, installed_pyarrow = _import_or_missing_pyarrow_option() + + assert installed_pyarrow + assert not isinstance(pyarrow, MissingPyArrow) diff --git a/test/unit/test_arrow_without_pandas.py b/test/unit/test_arrow_without_pandas.py new file mode 100644 index 0000000000..2951bc4493 --- /dev/null +++ b/test/unit/test_arrow_without_pandas.py @@ -0,0 +1,99 @@ +"""Regression tests for Arrow result fetching without pandas (the ``arrow`` extra). + +``snowflake.connector.options`` resolves its optional dependencies once at +first import, and pytest has already imported the connector, so the +pandas/pyarrow availability matrix cannot be exercised in-process. The probes +run in a subprocess where the unwanted package is made unimportable through a +shim on PYTHONPATH. + +The probes are marked ``pandas`` so they run in the CI environment that +installs both pandas and pyarrow; each subprocess then blocks the package it +needs absent. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap + +import pytest + + +def _run_probe(tmp_path, blocked_package: str, probe: str) -> None: + (tmp_path / f"{blocked_package}.py").write_text( + f"raise ImportError('{blocked_package} intentionally unavailable')\n", + encoding="utf-8", + ) + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [str(tmp_path), env.get("PYTHONPATH", "")] + ).rstrip(os.pathsep) + completed = subprocess.run( + [sys.executable, "-c", textwrap.dedent(probe)], + env=env, + capture_output=True, + text=True, + ) + assert completed.returncode == 0, completed.stderr + + +@pytest.mark.pandas +def test_arrow_support_does_not_require_pandas(tmp_path): + """With pandas unimportable, Arrow support must still be fully available.""" + _run_probe( + tmp_path, + "pandas", + """ + from types import SimpleNamespace + + from snowflake.connector import options + from snowflake.connector.errors import ProgrammingError + from snowflake.connector.result_batch import ArrowResultBatch + + assert options.installed_pyarrow is True + assert options.installed_pandas is False + + batch = SimpleNamespace( + _schema=[SimpleNamespace(name="VALUE", type_code=0)] + ) + table = ArrowResultBatch._create_empty_table(batch) + assert table.schema == options.pyarrow.schema( + [options.pyarrow.field("VALUE", options.pyarrow.int64())] + ) + + try: + ArrowResultBatch._check_can_use_pandas(batch) + except ProgrammingError: + pass + else: + raise AssertionError("pandas fetch support should remain unavailable") + """, + ) + + +@pytest.mark.pandas +def test_pandas_support_still_requires_pyarrow(tmp_path): + """With pandas importable but pyarrow missing, pandas features must fail fast.""" + _run_probe( + tmp_path, + "pyarrow", + """ + import pandas # noqa: F401 -- the scenario is pandas installed without pyarrow + + from snowflake.connector import options + from snowflake.connector.errors import MissingDependencyError + + assert options.installed_pyarrow is False + # pandas fetch APIs and write_pandas require pyarrow, so pandas support is + # reported unavailable and using the pandas module fails fast + assert options.installed_pandas is False + try: + options.pandas.DataFrame + except MissingDependencyError: + pass + else: + raise AssertionError("options.pandas should fail fast without pyarrow") + """, + )