snowflake-ml-python 1.4.1__py3-none-any.whl → 1.5.1__py3-none-any.whl

snowflake/ml/_internal/env_utils.py

@@ -13,10 +13,6 @@ from packaging import requirements, specifiers, utils as packaging_utils, versio
 
 import snowflake.connector
 from snowflake.ml._internal import env as snowml_env
-from snowflake.ml._internal.exceptions import (
-    error_codes,
-    exceptions as snowml_exceptions,
-)
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import context, exceptions, session
 from snowflake.snowpark._internal import utils as snowpark_utils
@@ -237,6 +233,72 @@ def get_local_installed_version_of_pip_package(pip_req: requirements.Requirement
     return new_pip_req
 
 
+class IncorrectLocalEnvironmentError(Exception):
+    ...
+
+
+def validate_local_installed_version_of_pip_package(pip_req: requirements.Requirement) -> None:
+    """Validate if the package is locally installed, and the local version meet the specifier of the requirements.
+
+    Args:
+        pip_req: A requirements.Requirement object showing the requirement.
+
+    Raises:
+        IncorrectLocalEnvironmentError: Raised when cannot find the local installation of the requested package.
+        IncorrectLocalEnvironmentError: Raised when the local installed version cannot meet the requirement.
+    """
+    try:
+        local_dist = importlib_metadata.distribution(pip_req.name)
+        local_dist_version = version.parse(local_dist.version)
+    except importlib_metadata.PackageNotFoundError:
+        raise IncorrectLocalEnvironmentError(f"Cannot find the local installation of the requested package {pip_req}.")
+
+    if not pip_req.specifier.contains(local_dist_version):
+        raise IncorrectLocalEnvironmentError(
+            f"The local installed version {local_dist_version} cannot meet the requirement {pip_req}."
+        )
+
+
+CONDA_PKG_NAME_TO_PYPI_MAP = {"pytorch": "torch"}
+
+
+def try_convert_conda_requirement_to_pip(conda_req: requirements.Requirement) -> requirements.Requirement:
+    """Return a new requirements.Requirement object whose name has been attempted to convert to name in pypi from conda.
+
+    Args:
+        conda_req: A requirements.Requirement object showing the requirement in conda.
+
+    Returns:
+        A new requirements.Requirement object showing the requirement in pypi.
+    """
+    pip_req = copy.deepcopy(conda_req)
+    pip_req.name = CONDA_PKG_NAME_TO_PYPI_MAP.get(conda_req.name, conda_req.name)
+    return pip_req
+
+
+def validate_py_runtime_version(provided_py_version_str: str) -> None:
+    """Validate the provided python version string with python version in current runtime.
+        If the major or minor is different, errors out.
+
+    Args:
+        provided_py_version_str: the provided python version string.
+
+    Raises:
+        IncorrectLocalEnvironmentError: Raised when the provided python version has different major or minor.
+    """
+    if provided_py_version_str != snowml_env.PYTHON_VERSION:
+        provided_py_version = version.parse(provided_py_version_str)
+        current_py_version = version.parse(snowml_env.PYTHON_VERSION)
+        if (
+            provided_py_version.major != current_py_version.major
+            or provided_py_version.minor != current_py_version.minor
+        ):
+            raise IncorrectLocalEnvironmentError(
+                f"Requested python version is {provided_py_version_str} "
+                f"while current Python version is {snowml_env.PYTHON_VERSION}. "
+            )
+
+
 def get_package_spec_with_supported_ops_only(req: requirements.Requirement) -> requirements.Requirement:
     """Get the package spec with supported ops only including ==, >=, <=, > and <
 
@@ -491,6 +553,9 @@ def load_conda_env_file(
         A tuple of Dict of conda dependencies after validated, optional pip requirements if exist
         and a string 'major.minor.patchlevel' of python version.
     """
+    if not path.exists():
+        return collections.defaultdict(list), None, None
+
     with open(path, encoding="utf-8") as f:
         env = yaml.safe_load(stream=f)
 
@@ -541,6 +606,9 @@ def load_requirements_file(path: pathlib.Path) -> List[requirements.Requirement]
     Returns:
         List of dependencies string after validated.
     """
+    if not path.exists():
+        return []
+
     with open(path, encoding="utf-8") as f:
         reqs = f.readlines()
 
@@ -568,33 +636,6 @@ def parse_python_version_string(dep: str) -> Optional[str]:
     return None
 
 
-def validate_py_runtime_version(provided_py_version_str: str) -> None:
-    """Validate the provided python version string with python version in current runtime.
-        If the major or minor is different, errors out.
-
-    Args:
-        provided_py_version_str: the provided python version string.
-
-    Raises:
-        SnowflakeMLException: Raised when the provided python version has different major or minor.
-    """
-    if provided_py_version_str != snowml_env.PYTHON_VERSION:
-        provided_py_version = version.parse(provided_py_version_str)
-        current_py_version = version.parse(snowml_env.PYTHON_VERSION)
-        if (
-            provided_py_version.major != current_py_version.major
-            or provided_py_version.minor != current_py_version.minor
-        ):
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.LOCAL_ENVIRONMENT_ERROR,
-                original_exception=RuntimeError(
-                    f"Unable to load model which is saved with Python {provided_py_version_str} "
-                    f"while current Python version is {snowml_env.PYTHON_VERSION}. "
-                    "To load model metadata only, set meta_only to True."
-                ),
-            )
-
-
 def _find_conda_dep_spec(
     conda_chan_deps: DefaultDict[str, List[requirements.Requirement]], pkg_name: str
 ) -> Optional[Tuple[str, requirements.Requirement]]:

snowflake/ml/_internal/exceptions/dataset_error_messages.py

@@ -0,0 +1,5 @@
+DATASET_ALREADY_EXISTS = "Dataset {} already exists."
+DATASET_VERSION_ALREADY_EXISTS = "Dataset {} version {} already exists."
+
+DATASET_NOT_EXIST = "Dataset {} does not exist or is inaccessible."
+DATASET_VERSION_NOT_EXIST = "Dataset {} version '{}' does not exist or is inaccessible."

snowflake/ml/_internal/exceptions/dataset_errors.py

@@ -0,0 +1,24 @@
+# Error code from Snowflake Python Connector.
+ERRNO_OBJECT_ALREADY_EXISTS = "002002"
+ERRNO_OBJECT_NOT_EXIST = "002043"
+ERRNO_FILES_ALREADY_EXISTING = "001030"
+ERRNO_VERSION_ALREADY_EXISTS = "092917"
+ERRNO_DATASET_NOT_EXIST = "399019"
+ERRNO_DATASET_VERSION_NOT_EXIST = "399012"
+ERRNO_DATASET_VERSION_ALREADY_EXISTS = "399020"
+
+
+class DatasetError(Exception):
+    """Base class for other exceptions."""
+
+
+class DatasetNotExistError(DatasetError):
+    """Raised when the requested Dataset does not exist."""
+
+
+class DatasetExistError(DatasetError):
+    """Raised when there is already an existing Dataset with the same name and version in selected schema."""
+
+
+class DatasetCannotDeleteError(DatasetError):
+    """Raised when a Dataset is unable to get deleted."""

snowflake/ml/_internal/exceptions/error_codes.py

@@ -105,3 +105,6 @@ UNFEASIBLE_ENVIRONMENT_ERROR = "2502"
 
 # Missing required client side dependency.
 CLIENT_DEPENDENCY_MISSING_ERROR = "2511"
+
+# Current client side snowpark-ml-python version is outdated and may have forward compatibility issue
+SNOWML_PACKAGE_OUTDATED = "2700"

snowflake/ml/_internal/lineage/data_source.py

@@ -0,0 +1,10 @@
+import dataclasses
+from typing import List, Optional
+
+
+@dataclasses.dataclass(frozen=True)
+class DataSource:
+    fully_qualified_name: str
+    version: str
+    url: str
+    exclude_cols: Optional[List[str]] = None

snowflake/ml/_internal/lineage/lineage_utils.py

@@ -0,0 +1,95 @@
+import copy
+import functools
+from typing import Any, Callable, List
+
+from snowflake import snowpark
+from snowflake.ml._internal.lineage import data_source
+
+DATA_SOURCES_ATTR = "_data_sources"
+
+
+def _get_datasources(*args: Any) -> List[data_source.DataSource]:
+    """Helper method for extracting data sources attribute from DataFrames in an argument list"""
+    result = []
+    for arg in args:
+        srcs = getattr(arg, DATA_SOURCES_ATTR, None)
+        if isinstance(srcs, list) and all(isinstance(s, data_source.DataSource) for s in srcs):
+            result += srcs
+    return result
+
+
+def _wrap_func(
+    fn: Callable[..., snowpark.DataFrame], data_sources: List[data_source.DataSource]
+) -> Callable[..., snowpark.DataFrame]:
+    """Wrap a DataFrame transform function to propagate data_sources to derived DataFrames."""
+
+    @functools.wraps(fn)
+    def wrapped(*args: Any, **kwargs: Any) -> snowpark.DataFrame:
+        df = fn(*args, **kwargs)
+        patch_dataframe(df, data_sources=data_sources, inplace=True)
+        return df
+
+    return wrapped
+
+
+def patch_dataframe(
+    df: snowpark.DataFrame, data_sources: List[data_source.DataSource], inplace: bool = False
+) -> snowpark.DataFrame:
+    """
+    Monkey patch a DataFrame to add attach the provided data_sources as an attribute of the DataFrame.
+    Also patches the DataFrame's transformation functions to propagate the new data sources attribute to
+    derived DataFrames.
+
+    Args:
+        df: DataFrame to be patched
+        data_sources: List of data sources for the DataFrame
+        inplace: If True, patches to DataFrame in-place. If False, creates a shallow copy of the DataFrame.
+
+    Returns:
+        Patched DataFrame
+    """
+    # Instance-level monkey-patches
+    funcs = [
+        "_with_plan",
+        "_lateral",
+        "group_by",
+        "group_by_grouping_sets",
+        "cube",
+        "pivot",
+        "rollup",
+        "cache_result",
+        "_to_df",  # RelationalGroupedDataFrame
+    ]
+    if not inplace:
+        df = copy.copy(df)
+    setattr(df, DATA_SOURCES_ATTR, data_sources)
+    for func in funcs:
+        fn = getattr(df, func, None)
+        if fn is not None:
+            setattr(df, func, _wrap_func(fn, data_sources=data_sources))
+    return df
+
+
+def _wrap_class_func(fn: Callable[..., snowpark.DataFrame]) -> Callable[..., snowpark.DataFrame]:
+    @functools.wraps(fn)
+    def wrapped(*args: Any, **kwargs: Any) -> snowpark.DataFrame:
+        df = fn(*args, **kwargs)
+        data_sources = _get_datasources(*args) + _get_datasources(*kwargs.values())
+        if data_sources:
+            patch_dataframe(df, data_sources, inplace=True)
+        return df
+
+    return wrapped
+
+
+# Class-level monkey-patches
+for klass, func_list in {
+    snowpark.DataFrame: [
+        "__copy__",
+    ],
+    snowpark.RelationalGroupedDataFrame: [],
+}.items():
+    assert isinstance(func_list, list)  # mypy
+    for func in func_list:
+        fn = getattr(klass, func)
+        setattr(klass, func, _wrap_class_func(fn))

snowflake/ml/_internal/telemetry.py

@@ -50,6 +50,7 @@ class TelemetryField(enum.Enum):
     # types of telemetry
     TYPE_FUNCTION_USAGE = "function_usage"
     TYPE_SNOWML_SPCS_USAGE = "snowml_spcs_usage"
+    TYPE_SNOWML_PIPELINE_USAGE = "snowml_pipeline_usage"
     # message keys for telemetry
     KEY_PROJECT = "project"
     KEY_SUBPROJECT = "subproject"

snowflake/ml/_internal/utils/identifier.py

@@ -156,7 +156,7 @@ def parse_schema_level_object_identifier(
     """
     res = _SF_SCHEMA_LEVEL_OBJECT_RE.fullmatch(path)
     if not res:
-        raise ValueError(f"Invalid identifier. It should start with database.schema.stage. Getting {path}")
+        raise ValueError(f"Invalid identifier. It should start with database.schema.object. Getting {path}")
     return (
         res.group("db"),
         res.group("schema"),

snowflake/ml/_internal/utils/sql_identifier.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional, Tuple
 
 from snowflake.ml._internal.utils import identifier
 
@@ -79,3 +79,16 @@ class SqlIdentifier(str):
 
 def to_sql_identifiers(list_of_str: List[str], *, case_sensitive: bool = False) -> List[SqlIdentifier]:
     return [SqlIdentifier(val, case_sensitive=case_sensitive) for val in list_of_str]
+
+
+def parse_fully_qualified_name(
+    name: str,
+) -> Tuple[Optional[SqlIdentifier], Optional[SqlIdentifier], SqlIdentifier]:
+    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 (
+        SqlIdentifier(db) if db else None,
+        SqlIdentifier(schema) if schema else None,
+        SqlIdentifier(object),
+    )

snowflake/ml/dataset/__init__.py

@@ -0,0 +1,11 @@
+from .dataset import Dataset, DatasetVersion
+from .dataset_factory import create_from_dataframe, load_dataset
+from .dataset_reader import DatasetReader
+
+__all__ = [
+    "Dataset",
+    "DatasetVersion",
+    "DatasetReader",
+    "create_from_dataframe",
+    "load_dataset",
+]