Optimize import time: Directly lookup target distributions instead of filtering manually (#2120)

Author: sfc-gh-rebeling

DESCRIPTION.md

@@ -11,6 +11,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Bumped pyOpenSSL dependency upper boundary from <25.0.0 to <26.0.0.
   - Removed the workaround for a Python 2.7 bug.
   - Added a <19.0.0 pin to pyarrow as a workaround to a bug affecting Azure Batch.
+  - Optimized distribution package lookup to speed up import.
 
 - v3.13.2(January 29, 2025)
   - Changed not to use scoped temporary objects.

src/snowflake/connector/options.py

@@ -7,7 +7,7 @@
 import importlib
 import os
 import warnings
-from importlib.metadata import distributions
+from importlib.metadata import PackageNotFoundError, distribution
 from logging import getLogger
 from types import ModuleType
 from typing import Union
@@ -85,13 +85,13 @@ def _import_or_missing_pandas_option() -> (
             os.environ["ARROW_DEFAULT_MEMORY_POOL"] = "system"
 
         # Check whether we have the currently supported pyarrow installed
-        installed_packages = {
-            package.metadata["Name"]: package for package in distributions()
-        }
-        if {"pyarrow", "snowflake-connector-python"} <= installed_packages.keys():
-            dependencies = installed_packages[
-                "snowflake-connector-python"
-            ].metadata.get_all("Requires-Dist", [])
+        try:
+            pyarrow_dist = distribution("pyarrow")
+            snowflake_connector_dist = distribution("snowflake-connector-python")
+
+            dependencies = snowflake_connector_dist.metadata.get_all(
+                "Requires-Dist", []
+            )
             pandas_pyarrow_extra = None
             for dependency in dependencies:
                 dep = Requirement(dependency)
@@ -103,16 +103,15 @@ def _import_or_missing_pandas_option() -> (
                     pandas_pyarrow_extra = dep
                     break
 
-            installed_pyarrow_version = installed_packages["pyarrow"].version
+            installed_pyarrow_version = pyarrow_dist.version
             if not pandas_pyarrow_extra.specifier.contains(installed_pyarrow_version):
                 warn_incompatible_dep(
                     "pyarrow", installed_pyarrow_version, pandas_pyarrow_extra
                 )
 
-        else:
+        except PackageNotFoundError as e:
             logger.info(
-                "Cannot determine if compatible pyarrow is installed because of missing package(s) from "
-                "{}".format(list(installed_packages.keys()))
+                f"Cannot determine if compatible pyarrow is installed because of missing package(s): {e}"
             )
         return pandas, pyarrow, True
     except ImportError:

test/integ/pandas/test_unit_options.py

@@ -18,7 +18,7 @@
     MissingPandas = None
     _import_or_missing_pandas_option = None
 
-from importlib.metadata import distributions
+from importlib.metadata import PackageNotFoundError, distribution
 
 
 @pytest.mark.skipif(
@@ -30,25 +30,23 @@ def test_pandas_option_reporting(caplog):
 
     This issue was brought to attention in: https://github.com/snowflakedb/snowflake-connector-python/issues/412
     """
-    modified_distributions = list(
-        d
-        for d in distributions()
-        if d.metadata["Name"]
-        not in (
-            "pyarrow",
-            "snowflake-connecctor-python",
-        )
-    )
+
+    def modified_distribution(name, *args, **kwargs):
+        if name in ["pyarrow", "snowflake-connector-python"]:
+            raise PackageNotFoundError("TestErrorMessage")
+        return distribution(name, *args, **kwargs)
+
     with mock.patch(
-        "snowflake.connector.options.distributions",
-        return_value=modified_distributions,
+        "snowflake.connector.options.distribution",
+        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)
         assert (
-            "Cannot determine if compatible pyarrow is installed because of missing package(s) "
-            "from "
-        ) in caplog.text
+            "Cannot determine if compatible pyarrow is installed because of missing package(s)"
+            in caplog.text
+        )
+        assert "TestErrorMessage" in caplog.text