snowflake-ml-python 1.6.2__tar.gz → 1.6.4__tar.gz

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/CHANGELOG.md

@@ -1,6 +1,31 @@
 # Release History
 
-## 1.6.2 (TBD)
+## 1.6.4
+
+### Bug Fixes
+
+- Registry: Fix an issue that leads to incident when using `ModelVersion.run` with service.
+
+## 1.6.3 (2024-10-07)
+
+- Model Registry (PrPr) has been removed.
+
+### Bug Fixes
+
+- Registry: Fix a bug that when package whose name does not follow PEP-508 is provided when logging the model,
+  an unexpected normalization is happening.
+- Registry: Fix `not a valid remote uri` error when logging mlflow models.
+- Registry: Fix a bug that `ModelVersion.run` is called in a nested way.
+- Registry: Fix an issue that leads to `log_model` failure when local package version contains parts other than
+  base version.
+
+### New Features
+
+- Data: Improve `DataConnector.to_pandas()` performance when loading from Snowpark DataFrames.
+- Model Registry: Allow users to set a model task while using `log_model`.
+- Feature Store: FeatureView supports ON_CREATE or ON_SCHEDULE initialize mode.
+
+## 1.6.2 (2024-09-04)
 
 ### Bug Fixes
 
@@ -18,8 +43,6 @@
 - Data: Add native batching support via `batch_size` and `drop_last_batch` arguments to `DataConnector.to_torch_dataset()`
 - Feature Store: update_feature_view() supports taking feature view object as argument.
 
-### Behavior Changes
-
 ## 1.6.1 (2024-08-12)
 
 ### Bug Fixes
@@ -42,8 +65,6 @@
 - Registry: Option to `enable_explainability` set to True by default for XGBoost, LightGBM and CatBoost as PuPr feature.
 - Registry: Option to `enable_explainability` when registering SHAP supported sklearn models.
 
-### Behavior Changes
-
 ## 1.6.0 (2024-07-29)
 
 ### Bug Fixes

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: snowflake-ml-python
-Version: 1.6.2
+Version: 1.6.4
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -247,7 +247,7 @@ Requires-Dist: pytimeparse<2,>=1.1.8
 Requires-Dist: pyyaml<7,>=6.0
 Requires-Dist: retrying<2,>=1.3.3
 Requires-Dist: s3fs<2024,>=2022.11
-Requires-Dist: scikit-learn<1.4,>=1.2.1
+Requires-Dist: scikit-learn<1.6,>=1.2.1
 Requires-Dist: scipy<2,>=1.9
 Requires-Dist: snowflake-connector-python[pandas]<4,>=3.5.0
 Requires-Dist: snowflake-snowpark-python<2,>=1.17.0
@@ -261,7 +261,7 @@ Requires-Dist: mlflow<2.4,>=2.1.0; extra == "all"
 Requires-Dist: peft<1,>=0.5.0; extra == "all"
 Requires-Dist: sentence-transformers<3,>=2.2.2; extra == "all"
 Requires-Dist: sentencepiece<1,>=0.1.95; extra == "all"
-Requires-Dist: shap==0.42.1; extra == "all"
+Requires-Dist: shap<1,>=0.42.0; extra == "all"
 Requires-Dist: tensorflow<3,>=2.10; extra == "all"
 Requires-Dist: tokenizers<1,>=0.10; extra == "all"
 Requires-Dist: torch<2.3.0,>=2.0.1; extra == "all"
@@ -276,7 +276,7 @@ Requires-Dist: peft<1,>=0.5.0; extra == "llm"
 Provides-Extra: mlflow
 Requires-Dist: mlflow<2.4,>=2.1.0; extra == "mlflow"
 Provides-Extra: shap
-Requires-Dist: shap==0.42.1; extra == "shap"
+Requires-Dist: shap<1,>=0.42.0; extra == "shap"
 Provides-Extra: tensorflow
 Requires-Dist: tensorflow<3,>=2.10; extra == "tensorflow"
 Provides-Extra: torch
@@ -373,7 +373,32 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 
 # Release History
 
-## 1.6.2 (TBD)
+## 1.6.4
+
+### Bug Fixes
+
+- Registry: Fix an issue that leads to incident when using `ModelVersion.run` with service.
+
+## 1.6.3 (2024-10-07)
+
+- Model Registry (PrPr) has been removed.
+
+### Bug Fixes
+
+- Registry: Fix a bug that when package whose name does not follow PEP-508 is provided when logging the model,
+  an unexpected normalization is happening.
+- Registry: Fix `not a valid remote uri` error when logging mlflow models.
+- Registry: Fix a bug that `ModelVersion.run` is called in a nested way.
+- Registry: Fix an issue that leads to `log_model` failure when local package version contains parts other than
+  base version.
+
+### New Features
+
+- Data: Improve `DataConnector.to_pandas()` performance when loading from Snowpark DataFrames.
+- Model Registry: Allow users to set a model task while using `log_model`.
+- Feature Store: FeatureView supports ON_CREATE or ON_SCHEDULE initialize mode.
+
+## 1.6.2 (2024-09-04)
 
 ### Bug Fixes
 
@@ -391,8 +416,6 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 - Data: Add native batching support via `batch_size` and `drop_last_batch` arguments to `DataConnector.to_torch_dataset()`
 - Feature Store: update_feature_view() supports taking feature view object as argument.
 
-### Behavior Changes
-
 ## 1.6.1 (2024-08-12)
 
 ### Bug Fixes
@@ -415,8 +438,6 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 - Registry: Option to `enable_explainability` set to True by default for XGBoost, LightGBM and CatBoost as PuPr feature.
 - Registry: Option to `enable_explainability` when registering SHAP supported sklearn models.
 
-### Behavior Changes
-
 ## 1.6.0 (2024-07-29)
 
 ### Bug Fixes

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/pyproject.toml

@@ -8,7 +8,7 @@ description = "The machine learning client library that is used for interacting
 classifiers = [ "Development Status :: 5 - Production/Stable", "Environment :: Console", "Environment :: Other Environment", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Information Technology", "Intended Audience :: System Administrators", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Topic :: Database", "Topic :: Software Development", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Application Frameworks", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Scientific/Engineering :: Information Analysis",]
 requires-python = ">=3.8, <3.12"
 dynamic = [ "version", "readme",]
-dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<4", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "fsspec[http]>=2022.11,<2024", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<2", "packaging>=20.9,<24", "pandas>=1.0.0,<3", "pyarrow", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2022.11,<2024", "scikit-learn>=1.2.1,<1.4", "scipy>=1.9,<2", "snowflake-connector-python[pandas]>=3.5.0,<4", "snowflake-snowpark-python>=1.17.0,<2", "sqlparse>=0.4,<1", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<2.1",]
+dependencies = [ "absl-py>=0.15,<2", "anyio>=3.5.0,<4", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "fsspec[http]>=2022.11,<2024", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<2", "packaging>=20.9,<24", "pandas>=1.0.0,<3", "pyarrow", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2022.11,<2024", "scikit-learn>=1.2.1,<1.6", "scipy>=1.9,<2", "snowflake-connector-python[pandas]>=3.5.0,<4", "snowflake-snowpark-python>=1.17.0,<2", "sqlparse>=0.4,<1", "typing-extensions>=4.1.0,<5", "xgboost>=1.7.3,<2.1",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"
@@ -24,12 +24,12 @@ Issues = "https://github.com/snowflakedb/snowflake-ml-python/issues"
 Changelog = "https://github.com/snowflakedb/snowflake-ml-python/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
-all = [ "catboost>=1.2.0, <2", "lightgbm>=3.3.5,<5", "mlflow>=2.1.0,<2.4", "peft>=0.5.0,<1", "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "shap==0.42.1", "tensorflow>=2.10,<3", "tokenizers>=0.10,<1", "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1", "transformers>=4.32.1,<5",]
+all = [ "catboost>=1.2.0, <2", "lightgbm>=3.3.5,<5", "mlflow>=2.1.0,<2.4", "peft>=0.5.0,<1", "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "shap>=0.42.0,<1", "tensorflow>=2.10,<3", "tokenizers>=0.10,<1", "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1", "transformers>=4.32.1,<5",]
 catboost = [ "catboost>=1.2.0, <2",]
 lightgbm = [ "lightgbm>=3.3.5,<5",]
 llm = [ "peft>=0.5.0,<1",]
 mlflow = [ "mlflow>=2.1.0,<2.4",]
-shap = [ "shap==0.42.1",]
+shap = [ "shap>=0.42.0,<1",]
 tensorflow = [ "tensorflow>=2.10,<3",]
 torch = [ "torch>=2.0.1,<2.3.0", "torchdata>=0.4,<1",]
 transformers = [ "sentence-transformers>=2.2.2,<3", "sentencepiece>=0.1.95,<1", "tokenizers>=0.10,<1", "transformers>=4.32.1,<5",]

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/cortex/__init__.py

@@ -1,5 +1,7 @@
 from snowflake.cortex._classify_text import ClassifyText
 from snowflake.cortex._complete import Complete, CompleteOptions
+from snowflake.cortex._embed_text_768 import EmbedText768
+from snowflake.cortex._embed_text_1024 import EmbedText1024
 from snowflake.cortex._extract_answer import ExtractAnswer
 from snowflake.cortex._sentiment import Sentiment
 from snowflake.cortex._summarize import Summarize
@@ -9,6 +11,8 @@ __all__ = [
     "ClassifyText",
     "Complete",
     "CompleteOptions",
+    "EmbedText768",
+    "EmbedText1024",
     "ExtractAnswer",
     "Sentiment",
     "Summarize",

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/cortex/_classify_text.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -33,4 +33,4 @@ def _classify_text_impl(
     categories: Union[List[str], snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, str_input, categories)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, str_input, categories))

snowflake_ml_python-1.6.4/snowflake/cortex/_embed_text_1024.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Union, cast
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def EmbedText1024(
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    """TextEmbed calls into the LLM inference service to embed the text.
+
+    Args:
+        model: A Column of strings representing the model to use for embedding. The value
+               of the strings must be within the SUPPORTED_MODELS list.
+        text: A Column of strings representing input text.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of vectors containing embeddings.
+    """
+
+    return _embed_text_1024_impl("snowflake.cortex.embed_text_1024", model, text, session=session)
+
+
+def _embed_text_1024_impl(
+    function: str,
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))

snowflake_ml_python-1.6.4/snowflake/cortex/_embed_text_768.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Union, cast
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def EmbedText768(
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    """TextEmbed calls into the LLM inference service to embed the text.
+
+    Args:
+        model: A Column of strings representing the model to use for embedding. The value
+               of the strings must be within the SUPPORTED_MODELS list.
+        text: A Column of strings representing input text.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of vectors containing embeddings.
+    """
+
+    return _embed_text_768_impl("snowflake.cortex.embed_text_768", model, text, session=session)
+
+
+def _embed_text_768_impl(
+    function: str,
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/cortex/_extract_answer.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -33,4 +33,4 @@ def _extract_answer_impl(
     question: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, from_text, question)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, from_text, question))

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/cortex/_sentiment.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -30,4 +30,4 @@ def _sentiment_impl(
     output = call_sql_function(function, session, text)
     if isinstance(output, snowpark.Column):
         return output
-    return float(output)
+    return float(cast(str, output))

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/cortex/_summarize.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -30,4 +30,4 @@ def _summarize_impl(
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text))

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/cortex/_translate.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -36,4 +36,4 @@ def _translate_impl(
     to_language: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text, from_language, to_language)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text, from_language, to_language))

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/cortex/_util.py

@@ -24,7 +24,7 @@ def call_sql_function(
     function: str,
     session: Optional[snowpark.Session],
     *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> Union[str, snowpark.Column]:
+) -> Union[str, List[float], snowpark.Column]:
     handle_as_column = False
 
     for arg in args:
@@ -32,9 +32,9 @@ def call_sql_function(
             handle_as_column = True
 
     if handle_as_column:
-        return cast(Union[str, snowpark.Column], _call_sql_function_column(function, *args))
+        return cast(Union[str, List[float], snowpark.Column], _call_sql_function_column(function, *args))
     return cast(
-        Union[str, snowpark.Column],
+        Union[str, List[float], snowpark.Column],
         _call_sql_function_immediate(function, session, *args),
     )
 
@@ -49,7 +49,7 @@ def _call_sql_function_immediate(
     function: str,
     session: Optional[snowpark.Session],
     *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> str:
+) -> Union[str, List[float]]:
     session = session or context.get_active_session()
     if session is None:
         raise SnowflakeAuthenticationException(

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/ml/_internal/env_utils.py

@@ -9,7 +9,7 @@ from importlib import metadata as importlib_metadata
 from typing import Any, DefaultDict, Dict, List, Optional, Tuple
 
 import yaml
-from packaging import requirements, specifiers, utils as packaging_utils, version
+from packaging import requirements, specifiers, version
 
 import snowflake.connector
 from snowflake.ml._internal import env as snowml_env
@@ -54,15 +54,12 @@ def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
     """
     try:
         r = requirements.Requirement(req_str)
-        r.name = packaging_utils.canonicalize_name(r.name)
 
         if r.name == "python":
             raise ValueError("Don't specify python as a dependency, use python version argument instead.")
     except requirements.InvalidRequirement:
         raise ValueError(f"Invalid package requirement {req_str} found.")
 
-    if r.marker:
-        raise ValueError("Markers is not supported in conda dependency.")
     return r
 
 
@@ -84,6 +81,8 @@ def _validate_conda_dependency_string(dep_str: str) -> Tuple[str, requirements.R
     channel_str, _, requirement_str = dep_str.rpartition("::")
     r = _validate_pip_requirement_string(requirement_str)
     if channel_str != "pip":
+        if r.marker:
+            raise ValueError("Markers is not supported in conda dependency.")
         if r.extras:
             raise ValueError("Extras is not supported in conda dependency.")
         if r.url:
@@ -221,7 +220,7 @@ def get_local_installed_version_of_pip_package(pip_req: requirements.Requirement
         else:
             return pip_req
     new_pip_req = copy.deepcopy(pip_req)
-    new_pip_req.specifier = specifiers.SpecifierSet(specifiers=f"=={local_dist_version}")
+    new_pip_req.specifier = specifiers.SpecifierSet(specifiers=f"=={version.parse(local_dist_version).base_version}")
     if not pip_req.specifier.contains(local_dist_version):
         warnings.warn(
             f"Package requirement {str(pip_req)} specified, while version {local_dist_version} is installed. "
@@ -513,6 +512,7 @@ def save_conda_env_file(
         )
 
     with open(path, "w", encoding="utf-8") as f:
+        yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
         yaml.safe_dump(env, stream=f, default_flow_style=False)
 
 

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/ml/_internal/exceptions/error_codes.py

@@ -58,6 +58,8 @@ NOT_FOUND = "2101"
 METHOD_NOT_ALLOWED = "2102"
 # Not implemented.
 NOT_IMPLEMENTED = "2103"
+# User needs to opt in to use a feature.
+OPT_IN_REQUIRED = "2104"
 
 # Calling an API with unsupported keywords/values.
 INVALID_ARGUMENT = "2110"

snowflake_ml_python-1.6.4/snowflake/ml/_internal/utils/db_utils.py

@@ -0,0 +1,50 @@
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.snowpark import session
+
+MAX_IDENTIFIER_LENGTH = 255
+
+
+class SnowflakeDbObjectType(Enum):
+    TABLE = "TABLE"
+    WAREHOUSE = "WAREHOUSE"
+
+
+def db_object_exists(
+    session: session.Session,
+    object_type: SnowflakeDbObjectType,
+    object_name: sql_identifier.SqlIdentifier,
+    *,
+    database_name: Optional[sql_identifier.SqlIdentifier] = None,
+    schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+    statement_params: Optional[Dict[str, Any]] = None,
+) -> bool:
+    """Check if object exists in database.
+
+    Args:
+        session: Active Snowpark Session.
+        object_type: Type of object to search for.
+        object_name: Name of object to search for.
+        database_name: Optional database name to search in. Only used if both schema is also provided.
+        schema_name: Optional schema to search in.
+        statement_params: Optional set of statement_params to include with queries.
+
+    Returns:
+        boolean indicating whether object exists.
+    """
+    optional_in_clause = ""
+    if database_name and schema_name:
+        optional_in_clause = f" IN {database_name}.{schema_name}"
+
+    result = (
+        query_result_checker.SqlResultValidator(
+            session,
+            f"""SHOW {object_type.value}S LIKE '{object_name}'{optional_in_clause}""",
+            statement_params=statement_params,
+        )
+        .has_column("name", allow_empty=True)  # TODO: Check this is actually what is returned from server
+        .validate()
+    )
+    return len(result) == 1

snowflake_ml_python-1.6.4/snowflake/ml/_internal/utils/service_logger.py

@@ -0,0 +1,63 @@
+import enum
+import logging
+import sys
+
+
+class LogColor(enum.Enum):
+    GREY = "\x1b[38;20m"
+    RED = "\x1b[31;20m"
+    BOLD_RED = "\x1b[31;1m"
+    YELLOW = "\x1b[33;20m"
+    BLUE = "\x1b[34;20m"
+    GREEN = "\x1b[32;20m"
+
+
+class CustomFormatter(logging.Formatter):
+
+    reset = "\x1b[0m"
+    log_format = "%(name)s [%(asctime)s] [%(levelname)s] %(message)s"
+
+    def __init__(self, info_color: LogColor) -> None:
+        super().__init__()
+        self.level_colors = {
+            logging.DEBUG: LogColor.GREY.value,
+            logging.INFO: info_color.value,
+            logging.WARNING: LogColor.YELLOW.value,
+            logging.ERROR: LogColor.RED.value,
+            logging.CRITICAL: LogColor.BOLD_RED.value,
+        }
+
+    def format(self, record: logging.LogRecord) -> str:
+        # default to DEBUG color
+        fmt = self.level_colors.get(record.levelno, self.level_colors[logging.DEBUG]) + self.log_format + self.reset
+        formatter = logging.Formatter(fmt)
+
+        # split the log message by lines and format each line individually
+        original_message = record.getMessage()
+        message_lines = original_message.splitlines()
+        formatted_lines = [
+            formatter.format(
+                logging.LogRecord(
+                    name=record.name,
+                    level=record.levelno,
+                    pathname=record.pathname,
+                    lineno=record.lineno,
+                    msg=line,
+                    args=None,
+                    exc_info=None,
+                )
+            )
+            for line in message_lines
+        ]
+
+        return "\n".join(formatted_lines)
+
+
+def get_logger(logger_name: str, info_color: LogColor) -> logging.Logger:
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(logging.INFO)
+    handler.setFormatter(CustomFormatter(info_color))
+    logger.addHandler(handler)
+    return logger

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/ml/_internal/utils/sql_identifier.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 from snowflake.ml._internal.utils import identifier
 
@@ -92,3 +92,27 @@ def parse_fully_qualified_name(
         SqlIdentifier(schema) if schema else None,
         SqlIdentifier(object),
     )
+
+
+def get_fully_qualified_name(
+    db: Union[SqlIdentifier, str, None],
+    schema: Union[SqlIdentifier, str, None],
+    object: Union[SqlIdentifier, str],
+    session_db: Optional[str] = None,
+    session_schema: Optional[str] = None,
+) -> str:
+    db_name: Optional[SqlIdentifier] = None
+    schema_name: Optional[SqlIdentifier] = None
+    if not db and session_db:
+        db_name = SqlIdentifier(session_db)
+    elif isinstance(db, str):
+        db_name = SqlIdentifier(db)
+    if not schema and session_schema:
+        schema_name = SqlIdentifier(session_schema)
+    elif isinstance(schema, str):
+        schema_name = SqlIdentifier(schema)
+    return identifier.get_schema_level_object_identifier(
+        db=db_name.identifier() if db_name else None,
+        schema=schema_name.identifier() if schema_name else None,
+        object_name=object.identifier() if isinstance(object, SqlIdentifier) else SqlIdentifier(object).identifier(),
+    )

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/ml/data/_internal/arrow_ingestor.py

@@ -11,7 +11,6 @@ import pyarrow as pa
 import pyarrow.dataset as pds
 
 from snowflake import snowpark
-from snowflake.connector import result_batch
 from snowflake.ml.data import data_ingestor, data_source, ingestor_utils
 
 _EMPTY_RECORD_BATCH = pa.RecordBatch.from_arrays([], [])
@@ -140,16 +139,7 @@ class ArrowIngestor(data_ingestor.DataIngestor):
                 #        We may be able to optimize this by splitting the result batches into
                 #        in-memory (first batch) and file URLs (subsequent batches) and creating a
                 #        union dataset.
-                result_batches = ingestor_utils.get_dataframe_result_batches(self._session, source)
-                sources.extend(
-                    b.to_arrow(self._session.connection)
-                    if isinstance(b, result_batch.ArrowResultBatch)
-                    else b.to_arrow()
-                    for b in result_batches
-                )
-                # HACK: Mitigate typing inconsistencies in Snowpark results
-                if len(sources) > 0:
-                    sources = [_cast_if_needed(s, sources[-1].schema) for s in sources]
+                sources.append(_cast_if_needed(ingestor_utils.get_dataframe_arrow_table(self._session, source)))
                 source_format = None  # Arrow Dataset expects "None" for in-memory datasets
             else:
                 raise RuntimeError(f"Unsupported data source type: {type(source)}")

{snowflake_ml_python-1.6.2 → snowflake_ml_python-1.6.4}/snowflake/ml/data/ingestor_utils.py

@@ -1,19 +1,17 @@
 from typing import List, Optional
 
 import fsspec
+import pyarrow as pa
 
 from snowflake import snowpark
-from snowflake.connector import result_batch
+from snowflake.connector import cursor as sf_cursor, result_batch
 from snowflake.ml.data import data_source
 from snowflake.ml.fileset import snowfs
 
 _TARGET_FILE_SIZE = 32 * 2**20  # The max file size for data loading.
 
 
-def get_dataframe_result_batches(
-    session: snowpark.Session, df_info: data_source.DataFrameInfo
-) -> List[result_batch.ResultBatch]:
-    """Retrieve the ResultBatches for a given query"""
+def _get_dataframe_cursor(session: snowpark.Session, df_info: data_source.DataFrameInfo) -> sf_cursor.SnowflakeCursor:
     cursor = session._conn._cursor
 
     if df_info.query_id:
@@ -29,12 +27,24 @@ def get_dataframe_result_batches(
     if cursor._prefetch_hook is None:
         raise RuntimeError("Loading data from result query failed unexpectedly. Please contact Snowflake support.")
     cursor._prefetch_hook()
+
+    return cursor
+
+
+def get_dataframe_result_batches(
+    session: snowpark.Session, df_info: data_source.DataFrameInfo
+) -> List[result_batch.ResultBatch]:
+    """Retrieve the ResultBatches for a given query"""
+    cursor = _get_dataframe_cursor(session, df_info)
     batches = cursor.get_result_batches()
-    if batches is None:
-        raise ValueError(
-            "Failed to retrieve training data. Query status:" f" {session._conn._conn.get_query_status(query_id)}"
-        )
-    return batches
+    return batches or []
+
+
+def get_dataframe_arrow_table(session: snowpark.Session, df_info: data_source.DataFrameInfo) -> pa.Table:
+    """Retrieve the full in-memory result for a given query"""
+    cursor = _get_dataframe_cursor(session, df_info)
+    table = cursor.fetch_arrow_all()  # type: ignore[call-overload]
+    return table
 
 
 def get_dataset_filesystem(