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

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

@@ -1,6 +1,26 @@
 # Release History
 
-## 1.6.1 (TBD)
+## 1.6.2 (TBD)
+
+### Bug Fixes
+
+- Modeling: Support XGBoost version that is larger than 2.
+
+- Data: Fix multiple epoch iteration over `DataConnector.to_torch_datapipe()` DataPipes.
+- Generic: Fix a bug that when an invalid name is provided to argument where fully qualified name is expected, it will
+  be parsed wrongly. Now it raises an exception correctly.
+- Model Explainability: Handle explanations for multiclass XGBoost classification models
+- Model Explainability: Workarounds and better error handling for XGB>2.1.0 not working with SHAP==0.42.1
+
+### New Features
+
+- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
+- 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
 
@@ -17,7 +37,6 @@
 ### New Features
 
 - Enable `set_params` to set the parameters of the underlying sklearn estimator, if the snowflake-ml model has been fit.
-- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
 - Data: Add `snowflake.ml.data.ingestor_utils` module with utility functions helpful for `DataIngestor` implementations.
 - Data: Add new `to_torch_dataset()` connector to `DataConnector` to replace deprecated DataPipe.
 - Registry: Option to `enable_explainability` set to True by default for XGBoost, LightGBM and CatBoost as PuPr feature.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: snowflake-ml-python
-Version: 1.6.1
+Version: 1.6.2
 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: 
@@ -253,7 +253,7 @@ Requires-Dist: snowflake-connector-python[pandas]<4,>=3.5.0
 Requires-Dist: snowflake-snowpark-python<2,>=1.17.0
 Requires-Dist: sqlparse<1,>=0.4
 Requires-Dist: typing-extensions<5,>=4.1.0
-Requires-Dist: xgboost<2,>=1.7.3
+Requires-Dist: xgboost<2.1,>=1.7.3
 Provides-Extra: all
 Requires-Dist: catboost<2,>=1.2.0; extra == "all"
 Requires-Dist: lightgbm<5,>=3.3.5; extra == "all"
@@ -373,7 +373,27 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 
 # Release History
 
-## 1.6.1 (TBD)
+## 1.6.2 (TBD)
+
+### Bug Fixes
+
+- Modeling: Support XGBoost version that is larger than 2.
+
+- Data: Fix multiple epoch iteration over `DataConnector.to_torch_datapipe()` DataPipes.
+- Generic: Fix a bug that when an invalid name is provided to argument where fully qualified name is expected, it will
+  be parsed wrongly. Now it raises an exception correctly.
+- Model Explainability: Handle explanations for multiclass XGBoost classification models
+- Model Explainability: Workarounds and better error handling for XGB>2.1.0 not working with SHAP==0.42.1
+
+### New Features
+
+- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
+- 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
 
@@ -390,7 +410,6 @@ be compatibility issues. Server-side functionality that `snowflake-ml-python` de
 ### New Features
 
 - Enable `set_params` to set the parameters of the underlying sklearn estimator, if the snowflake-ml model has been fit.
-- Data: Add top-level exports for `DataConnector` and `DataSource` to `snowflake.ml.data`.
 - Data: Add `snowflake.ml.data.ingestor_utils` module with utility functions helpful for `DataIngestor` implementations.
 - Data: Add new `to_torch_dataset()` connector to `DataConnector` to replace deprecated DataPipe.
 - Registry: Option to `enable_explainability` set to True by default for XGBoost, LightGBM and CatBoost as PuPr feature.

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.2}/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",]
+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",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/telemetry.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import contextvars
 import enum
 import functools
 import inspect
@@ -12,6 +13,7 @@ from typing import (
     List,
     Mapping,
     Optional,
+    Set,
     Tuple,
     TypeVar,
     Union,
@@ -28,7 +30,7 @@ from snowflake.ml._internal.exceptions import (
     exceptions as snowml_exceptions,
 )
 from snowflake.snowpark import dataframe, exceptions as snowpark_exceptions, session
-from snowflake.snowpark._internal import utils
+from snowflake.snowpark._internal import server_connection, utils
 
 _log_counter = 0
 _FLUSH_SIZE = 10
@@ -85,6 +87,122 @@ class TelemetryField(enum.Enum):
     FUNC_CAT_USAGE = "usage"
 
 
+class _TelemetrySourceType(enum.Enum):
+    # Automatically inferred telemetry/statement parameters
+    AUTO_TELEMETRY = "SNOWML_AUTO_TELEMETRY"
+    # Mixture of manual and automatic telemetry/statement parameters
+    AUGMENT_TELEMETRY = "SNOWML_AUGMENT_TELEMETRY"
+
+
+_statement_params_context_var: contextvars.ContextVar[Dict[str, str]] = contextvars.ContextVar("statement_params")
+
+
+class _StatementParamsPatchManager:
+    def __init__(self) -> None:
+        self._patch_cache: Set[server_connection.ServerConnection] = set()
+        self._context_var: contextvars.ContextVar[Dict[str, str]] = _statement_params_context_var
+
+    def apply_patches(self) -> None:
+        try:
+            # Apply patching to all active sessions in case of multiple
+            for sess in session._get_active_sessions():
+                # Check patch cache here to avoid unnecessary context switches
+                if self._get_target(sess) not in self._patch_cache:
+                    self._patch_session(sess)
+        except snowpark_exceptions.SnowparkSessionException:
+            pass
+
+    def set_statement_params(self, statement_params: Dict[str, str]) -> None:
+        # Only set value if not already set in context
+        if not self._context_var.get({}):
+            self._context_var.set(statement_params)
+
+    def _get_target(self, session: session.Session) -> server_connection.ServerConnection:
+        return cast(server_connection.ServerConnection, session._conn)
+
+    def _patch_session(self, session: session.Session, throw_on_patch_fail: bool = False) -> None:
+        # Extract target
+        try:
+            target = self._get_target(session)
+        except AttributeError:
+            if throw_on_patch_fail:
+                raise
+            # TODO: Log a warning, this probably means there was a breaking change in Snowpark/SnowflakeConnection
+            return
+
+        # Check if session has already been patched
+        if target in self._patch_cache:
+            return
+        self._patch_cache.add(target)
+
+        functions = [
+            ("execute_and_notify_query_listener", "_statement_params"),
+            ("execute_async_and_notify_query_listener", "_statement_params"),
+        ]
+
+        for func, param_name in functions:
+            try:
+                self._patch_with_statement_params(target, func, param_name=param_name)
+            except AttributeError:
+                if throw_on_patch_fail:  # primarily used for testing
+                    raise
+                # TODO: Log a warning, this probably means there was a breaking change in Snowpark/SnowflakeConnection
+                pass
+
+    def _patch_with_statement_params(
+        self, target: object, function_name: str, param_name: str = "statement_params"
+    ) -> None:
+        func = getattr(target, function_name)
+        assert callable(func)
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Retrieve context level statement parameters
+            context_params = self._context_var.get(dict())
+            if not context_params:
+                # Exit early if not in SnowML (decorator) context
+                return func(*args, **kwargs)
+
+            # Extract any explicitly provided statement parameters
+            orig_kwargs = dict(kwargs)
+            in_params = kwargs.pop(param_name, None) or {}
+
+            # Inject a special flag to statement parameters so we can filter out these patched logs if necessary
+            # Calls that include SnowML telemetry are tagged with "SNOWML_AUGMENT_TELEMETRY"
+            # and calls without SnowML telemetry are tagged with "SNOWML_AUTO_TELEMETRY"
+            if TelemetryField.KEY_PROJECT.value in in_params:
+                context_params["snowml_telemetry_type"] = _TelemetrySourceType.AUGMENT_TELEMETRY.value
+            else:
+                context_params["snowml_telemetry_type"] = _TelemetrySourceType.AUTO_TELEMETRY.value
+
+            # Apply any explicitly provided statement parameters and result into function call
+            context_params.update(in_params)
+            kwargs[param_name] = context_params
+
+            try:
+                return func(*args, **kwargs)
+            except TypeError as e:
+                if str(e).endswith(f"unexpected keyword argument '{param_name}'"):
+                    # TODO: Log warning that this patch is invalid
+                    # Unwrap function for future invocations
+                    setattr(target, function_name, func)
+                    return func(*args, **orig_kwargs)
+                else:
+                    raise
+
+        setattr(target, function_name, wrapper)
+
+    def __getstate__(self) -> Dict[str, Any]:
+        return {}
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        # unpickling does not call __init__ by default, do it manually here
+        self.__init__()  # type: ignore[misc]
+
+
+_patch_manager = _StatementParamsPatchManager()
+
+
 def get_statement_params(
     project: str, subproject: Optional[str] = None, class_name: Optional[str] = None
 ) -> Dict[str, Any]:
@@ -375,7 +493,18 @@ def send_api_usage_telemetry(
                         obj._statement_params = statement_params  # type: ignore[assignment]
                 return obj
 
+            # Set up framework-level credit usage instrumentation
+            ctx = contextvars.copy_context()
+            _patch_manager.apply_patches()
+
+            # This function should be executed with ctx.run()
+            def execute_func_with_statement_params() -> _ReturnValue:
+                _patch_manager.set_statement_params(statement_params)
+                result = func(*args, **kwargs)
+                return update_stmt_params_if_snowpark_df(result, statement_params)
+
             # prioritize `conn_attr_name` over the active session
+            telemetry_enabled = True
             if conn_attr_name:
                 # raise AttributeError if conn attribute does not exist in `self`
                 conn = operator.attrgetter(conn_attr_name)(args[0])
@@ -387,22 +516,17 @@ def send_api_usage_telemetry(
             else:
                 try:
                     active_session = next(iter(session._get_active_sessions()))
-                # server no default session
+                    conn = active_session._conn._conn
+                    telemetry_enabled = active_session.telemetry_enabled
                 except snowpark_exceptions.SnowparkSessionException:
-                    try:
-                        return update_stmt_params_if_snowpark_df(func(*args, **kwargs), statement_params)
-                    except Exception as e:
-                        if isinstance(e, snowml_exceptions.SnowflakeMLException):
-                            raise e.original_exception.with_traceback(e.__traceback__) from None
-                        # suppress SnowparkSessionException from telemetry in the stack trace
-                        raise e from None
-
-                conn = active_session._conn._conn
-                if (not active_session.telemetry_enabled) or (conn is None):
-                    try:
-                        return update_stmt_params_if_snowpark_df(func(*args, **kwargs), statement_params)
-                    except snowml_exceptions.SnowflakeMLException as e:
-                        raise e.original_exception from e
+                    conn = None
+
+            if conn is None or not telemetry_enabled:
+                # Telemetry not enabled, just execute without our additional telemetry logic
+                try:
+                    return ctx.run(execute_func_with_statement_params)
+                except snowml_exceptions.SnowflakeMLException as e:
+                    raise e.original_exception from e
 
             # TODO(hayu): [SNOW-750287] Optimize telemetry client to a singleton.
             telemetry = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject_name)
@@ -415,11 +539,11 @@ def send_api_usage_telemetry(
                 custom_tags=custom_tags,
             )
             try:
-                res = func(*args, **kwargs)
+                return ctx.run(execute_func_with_statement_params)
             except Exception as e:
                 if not isinstance(e, snowml_exceptions.SnowflakeMLException):
                     # already handled via a nested decorated function
-                    if hasattr(e, "_snowflake_ml_handled") and e._snowflake_ml_handled:
+                    if getattr(e, "_snowflake_ml_handled", False):
                         raise e
                     if isinstance(e, snowpark_exceptions.SnowparkClientException):
                         me = snowml_exceptions.SnowflakeMLException(
@@ -438,8 +562,6 @@ def send_api_usage_telemetry(
                     raise me.original_exception from None
                 else:
                     raise me.original_exception from e
-            else:
-                return update_stmt_params_if_snowpark_df(res, statement_params)
             finally:
                 telemetry.send_function_usage_telemetry(**telemetry_args)
                 global _log_counter

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/utils/identifier.py

@@ -10,9 +10,11 @@ SF_QUOTED_IDENTIFIER = '"(?:[^"]|"")*"'
 _SF_IDENTIFIER = f"({_SF_UNQUOTED_CASE_INSENSITIVE_IDENTIFIER}|{SF_QUOTED_IDENTIFIER})"
 SF_IDENTIFIER_RE = re.compile(_SF_IDENTIFIER)
 _SF_SCHEMA_LEVEL_OBJECT = (
-    rf"(?:(?:(?P<db>{_SF_IDENTIFIER})\.)?(?P<schema>{_SF_IDENTIFIER})\.)?(?P<object>{_SF_IDENTIFIER})(?P<others>.*)"
+    rf"(?:(?:(?P<db>{_SF_IDENTIFIER})\.)?(?P<schema>{_SF_IDENTIFIER})\.)?(?P<object>{_SF_IDENTIFIER})"
 )
+_SF_STAGE_PATH = rf"{_SF_SCHEMA_LEVEL_OBJECT}(?P<path>.*)"
 _SF_SCHEMA_LEVEL_OBJECT_RE = re.compile(_SF_SCHEMA_LEVEL_OBJECT)
+_SF_STAGE_PATH_RE = re.compile(_SF_STAGE_PATH)
 
 UNQUOTED_CASE_INSENSITIVE_RE = re.compile(f"^({_SF_UNQUOTED_CASE_INSENSITIVE_IDENTIFIER})$")
 UNQUOTED_CASE_SENSITIVE_RE = re.compile(f"^({_SF_UNQUOTED_CASE_SENSITIVE_IDENTIFIER})$")
@@ -139,29 +141,61 @@ def rename_to_valid_snowflake_identifier(name: str) -> str:
 
 
 def parse_schema_level_object_identifier(
+    object_name: str,
+) -> Tuple[Union[str, Any], Union[str, Any], Union[str, Any]]:
+    """Parse a string which starts with schema level object.
+
+    Args:
+        object_name: A string starts with a schema level object path, which is in the format
+            '<db>.<schema>.<object_name>'. Here, '<db>', '<schema>' and '<object_name>' are all snowflake identifiers.
+
+    Returns:
+        A tuple of 3 strings in the form of (db, schema, object_name).
+
+    Raises:
+        ValueError: If the id is invalid.
+    """
+    res = _SF_SCHEMA_LEVEL_OBJECT_RE.fullmatch(object_name)
+    if not res:
+        raise ValueError(
+            "Invalid identifier because it does not follow the pattern. "
+            f"It should start with [[database.]schema.]object. Getting {object_name}"
+        )
+    return (
+        res.group("db"),
+        res.group("schema"),
+        res.group("object"),
+    )
+
+
+def parse_snowflake_stage_path(
     path: str,
 ) -> Tuple[Union[str, Any], Union[str, Any], Union[str, Any], Union[str, Any]]:
-    """Parse a string which starts with schema level object.
+    """Parse a string which represents a snowflake stage path.
 
     Args:
-        path: A string starts with a schema level object path, which is in the format '<db>.<schema>.<object_name>'.
-            Here, '<db>', '<schema>' and '<object_name>' are all snowflake identifiers.
+        path: A string starts with a schema level object path, which is in the format
+            '<db>.<schema>.<object_name><path>'. Here, '<db>', '<schema>' and '<object_name>' are all snowflake
+            identifiers.
 
     Returns:
-        A tuple of 4 strings in the form of (db, schema, object_name, others). 'db', 'schema', 'object_name' are parsed
-            from the schema level object and 'others' are all the content post to the object.
+        A tuple of 4 strings in the form of (db, schema, object_name, path). 'db', 'schema', 'object_name' are parsed
+            from the schema level object and 'path' are all the content post to the object.
 
     Raises:
         ValueError: If the id is invalid.
     """
-    res = _SF_SCHEMA_LEVEL_OBJECT_RE.fullmatch(path)
+    res = _SF_STAGE_PATH_RE.fullmatch(path)
     if not res:
-        raise ValueError(f"Invalid identifier. It should start with database.schema.object. Getting {path}")
+        raise ValueError(
+            "Invalid identifier because it does not follow the pattern. "
+            f"It should start with [[database.]schema.]object. Getting {path}"
+        )
     return (
         res.group("db"),
         res.group("schema"),
         res.group("object"),
-        res.group("others"),
+        res.group("path"),
     )
 
 
@@ -175,8 +209,11 @@ def is_fully_qualified_name(name: str) -> bool:
     Returns:
         bool: True if the name is fully qualified, False otherwise.
     """
-    res = parse_schema_level_object_identifier(name)
-    return res[0] is not None and res[1] is not None and res[2] is not None and not res[3]
+    try:
+        res = parse_schema_level_object_identifier(name)
+        return all(res)
+    except ValueError:
+        return False
 
 
 def get_schema_level_object_identifier(

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/utils/snowflake_env.py

@@ -2,7 +2,7 @@ import enum
 from typing import Any, Dict, Optional, TypedDict, cast
 
 from packaging import version
-from typing_extensions import Required
+from typing_extensions import NotRequired, Required
 
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import session
@@ -52,7 +52,7 @@ class SnowflakeCloudType(enum.Enum):
 
 
 class SnowflakeRegion(TypedDict):
-    region_group: Required[str]
+    region_group: NotRequired[str]
     snowflake_region: Required[str]
     cloud: Required[SnowflakeCloudType]
     region: Required[str]
@@ -64,23 +64,33 @@ def get_regions(
 ) -> Dict[str, SnowflakeRegion]:
     res = (
         query_result_checker.SqlResultValidator(sess, "SHOW REGIONS", statement_params=statement_params)
-        .has_column("region_group")
         .has_column("snowflake_region")
         .has_column("cloud")
         .has_column("region")
         .has_column("display_name")
         .validate()
     )
-    return {
-        f"{r.region_group}.{r.snowflake_region}": SnowflakeRegion(
-            region_group=r.region_group,
-            snowflake_region=r.snowflake_region,
-            cloud=SnowflakeCloudType.from_value(r.cloud),
-            region=r.region,
-            display_name=r.display_name,
-        )
-        for r in res
-    }
+    res_dict = {}
+    for r in res:
+        if hasattr(r, "region_group") and r.region_group:
+            key = f"{r.region_group}.{r.snowflake_region}"
+            res_dict[key] = SnowflakeRegion(
+                region_group=r.region_group,
+                snowflake_region=r.snowflake_region,
+                cloud=SnowflakeCloudType.from_value(r.cloud),
+                region=r.region,
+                display_name=r.display_name,
+            )
+        else:
+            key = r.snowflake_region
+            res_dict[key] = SnowflakeRegion(
+                snowflake_region=r.snowflake_region,
+                cloud=SnowflakeCloudType.from_value(r.cloud),
+                region=r.region,
+                display_name=r.display_name,
+            )
+
+    return res_dict
 
 
 def get_current_region_id(sess: session.Session, *, statement_params: Optional[Dict[str, Any]] = None) -> str:

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

@@ -84,7 +84,7 @@ def to_sql_identifiers(list_of_str: List[str], *, case_sensitive: bool = False)
 def parse_fully_qualified_name(
     name: str,
 ) -> Tuple[Optional[SqlIdentifier], Optional[SqlIdentifier], SqlIdentifier]:
-    db, schema, object, _ = identifier.parse_schema_level_object_identifier(name)
+    db, schema, object = identifier.parse_schema_level_object_identifier(name)
 
     assert name is not None, f"Unable parse the input name `{name}` as fully qualified."
     return (

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/utils/table_manager.py

@@ -1,7 +1,8 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 from snowflake import snowpark
-from snowflake.ml._internal.utils import formatting, query_result_checker
+from snowflake.ml._internal.utils import formatting, identifier, query_result_checker
+from snowflake.snowpark import types
 
 """Table_manager is a set of utils that helps create tables.
 
@@ -104,3 +105,20 @@ def get_table_schema(session: snowpark.Session, table_name: str, qualified_schem
     for row in result:
         schema_dict[row["name"]] = row["type"]
     return schema_dict
+
+
+def get_table_schema_types(
+    session: snowpark.Session,
+    database: str,
+    schema: str,
+    table_name: str,
+) -> Dict[str, types.DataType]:
+    fully_qualified_table_name = identifier.get_schema_level_object_identifier(
+        db=database, schema=schema, object_name=table_name
+    )
+    struct_fields: List[types.StructField] = session.table(fully_qualified_table_name).schema.fields
+
+    schema_dict: Dict[str, types.DataType] = {}
+    for field in struct_fields:
+        schema_dict[field.name] = field.datatype
+    return schema_dict

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.2}/snowflake/ml/_internal/utils/uri.py

@@ -53,7 +53,7 @@ def get_uri_scheme(uri: str) -> str:
 def get_uri_from_snowflake_stage_path(stage_path: str) -> str:
     """Generates a URI from Snowflake stage path."""
     assert stage_path.startswith("@")
-    (db, schema, stage, path) = identifier.parse_schema_level_object_identifier(
+    (db, schema, stage, path) = identifier.parse_snowflake_stage_path(
         posixpath.normpath(identifier.remove_prefix(stage_path, "@"))
     )
     return urlunparse(
@@ -70,7 +70,7 @@ def get_uri_from_snowflake_stage_path(stage_path: str) -> str:
 
 def get_stage_and_path(stage_path: str) -> Tuple[str, str]:
     assert stage_path.startswith("@"), f"stage path should start with @, actual: {stage_path}"
-    (db, schema, stage, path) = identifier.parse_schema_level_object_identifier(
+    (db, schema, stage, path) = identifier.parse_snowflake_stage_path(
         posixpath.normpath(identifier.remove_prefix(stage_path, "@"))
     )
     full_qualified_stage = "@" + identifier.get_schema_level_object_identifier(db, schema, stage)

{snowflake_ml_python-1.6.1 → snowflake_ml_python-1.6.2}/snowflake/ml/data/data_connector.py

@@ -1,3 +1,4 @@
+import os
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Type, TypeVar
 
 import numpy.typing as npt
@@ -7,6 +8,10 @@ from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml.data import data_ingestor, data_source
 from snowflake.ml.data._internal.arrow_ingestor import ArrowIngestor
+from snowflake.ml.modeling._internal.constants import (
+    IN_ML_RUNTIME_ENV_VAR,
+    USE_OPTIMIZED_DATA_INGESTOR,
+)
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -142,32 +147,41 @@ class DataConnector:
         Returns:
             A Pytorch iterable datapipe that yield data.
         """
-        from torch.utils.data.datapipes import iter as torch_iter
+        from snowflake.ml.data import torch_utils
 
-        return torch_iter.IterableWrapper(  # type: ignore[no-untyped-call]
-            self._ingestor.to_batches(batch_size, shuffle, drop_last_batch)
+        return torch_utils.TorchDataPipeWrapper(
+            self._ingestor, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last_batch
         )
 
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
         subproject_extractor=lambda self: type(self).__name__,
-        func_params_to_log=["shuffle"],
+        func_params_to_log=["batch_size", "shuffle", "drop_last_batch"],
     )
-    def to_torch_dataset(self, *, shuffle: bool = False) -> "torch_data.IterableDataset":  # type: ignore[type-arg]
+    def to_torch_dataset(
+        self, *, batch_size: int = 1, shuffle: bool = False, drop_last_batch: bool = True
+    ) -> "torch_data.IterableDataset":  # type: ignore[type-arg]
         """Transform the Snowflake data into a PyTorch Iterable Dataset to be used with a DataLoader.
 
         Return a PyTorch Dataset which iterates on rows of data.
 
         Args:
+            batch_size: It specifies the size of each data batch which will be yielded in the result dataset.
+                Batching is pushed down to data ingestion level which may be more performant than DataLoader
+                batching.
             shuffle: It specifies whether the data will be shuffled. If True, files will be shuffled, and
                 rows in each file will also be shuffled.
+            drop_last_batch: Whether the last batch of data should be dropped. If set to be true,
+                then the last batch will get dropped if its size is smaller than the given batch_size.
 
         Returns:
             A PyTorch Iterable Dataset that yields data.
         """
-        from snowflake.ml.data import torch_dataset
+        from snowflake.ml.data import torch_utils
 
-        return torch_dataset.TorchDataset(self._ingestor, shuffle)
+        return torch_utils.TorchDatasetWrapper(
+            self._ingestor, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last_batch
+        )
 
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
@@ -184,3 +198,15 @@ class DataConnector:
             A Pandas DataFrame.
         """
         return self._ingestor.to_pandas(limit)
+
+
+# Switch to use Runtime's Data Ingester if running in ML runtime
+# Fail silently if the data ingester is not found
+if os.getenv(IN_ML_RUNTIME_ENV_VAR) and os.getenv(USE_OPTIMIZED_DATA_INGESTOR):
+    try:
+        from runtime_external_entities import get_ingester_class
+
+        DataConnector.DEFAULT_INGESTOR_CLASS = get_ingester_class()
+    except ImportError:
+        """Runtime Default Ingester not found, ignore"""
+        pass