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

snowflake/ml/feature_store/feature_view.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import json
+import re
 from collections import OrderedDict
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -26,21 +27,45 @@ from snowflake.snowpark.types import (
 )
 
 _FEATURE_VIEW_NAME_DELIMITER = "$"
-_TIMESTAMP_COL_PLACEHOLDER = "FS_TIMESTAMP_COL_PLACEHOLDER_VAL"
+_LEGACY_TIMESTAMP_COL_PLACEHOLDER_VALS = ["FS_TIMESTAMP_COL_PLACEHOLDER_VAL", "NULL"]
+_TIMESTAMP_COL_PLACEHOLDER = "NULL"
 _FEATURE_OBJ_TYPE = "FEATURE_OBJ_TYPE"
+# Feature view version rule is aligned with dataset version rule in SQL.
+_FEATURE_VIEW_VERSION_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_.\-]*$")
+_FEATURE_VIEW_VERSION_MAX_LENGTH = 128
 
 
-class FeatureViewVersion(SqlIdentifier):
+@dataclass(frozen=True)
+class _FeatureViewMetadata:
+    """Represent metadata tracked on top of FV backend object"""
+
+    entities: List[str]
+    timestamp_col: str
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self))
+
+    @classmethod
+    def from_json(cls, json_str: str) -> _FeatureViewMetadata:
+        state_dict = json.loads(json_str)
+        return cls(**state_dict)
+
+
+class FeatureViewVersion(str):
     def __new__(cls, version: str) -> FeatureViewVersion:
-        if _FEATURE_VIEW_NAME_DELIMITER in version:
+        if not _FEATURE_VIEW_VERSION_RE.match(version) or len(version) > _FEATURE_VIEW_VERSION_MAX_LENGTH:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(f"{_FEATURE_VIEW_NAME_DELIMITER} is not allowed in version: {version}."),
+                original_exception=ValueError(
+                    f"`{version}` is not a valid feature view version. "
+                    "It must start with letter or digit, and followed by letter, digit, '_', '-' or '.'. "
+                    f"The length limit is {_FEATURE_VIEW_VERSION_MAX_LENGTH}."
+                ),
             )
-        return super().__new__(cls, version)  # type: ignore[return-value]
+        return super().__new__(cls, version)
 
     def __init__(self, version: str) -> None:
-        super().__init__(version)
+        super().__init__()
 
 
 class FeatureViewStatus(Enum):
@@ -97,12 +122,13 @@ class FeatureView:
         timestamp_col: Optional[str] = None,
         refresh_freq: Optional[str] = None,
         desc: str = "",
+        **_kwargs: Any,
     ) -> None:
         """
         Create a FeatureView instance.
 
         Args:
-            name: name of the FeatureView. NOTE: FeatureView name will be capitalized.
+            name: name of the FeatureView. NOTE: following Snowflake identifier rule
             entities: entities that the FeatureView is associated with.
             feature_df: Snowpark DataFrame containing data source and all feature feature_df logics.
                 Final projection of the DataFrame should contain feature names, join keys and timestamp(if applicable).
@@ -116,6 +142,7 @@ class FeatureView:
                 NOTE: If refresh_freq is not provided, then FeatureView will be registered as View on Snowflake backend
                     and there won't be extra storage cost.
             desc: description of the FeatureView.
+            _kwargs: reserved kwargs for system generated args. NOTE: DO NOT USE.
         """
 
         self._name: SqlIdentifier = SqlIdentifier(name)
@@ -125,6 +152,7 @@ class FeatureView:
             SqlIdentifier(timestamp_col) if timestamp_col is not None else None
         )
         self._desc: str = desc
+        self._infer_schema_df: DataFrame = _kwargs.get("_infer_schema_df", self._feature_df)
         self._query: str = self._get_query()
         self._version: Optional[FeatureViewVersion] = None
         self._status: FeatureViewStatus = FeatureViewStatus.DRAFT
@@ -271,7 +299,7 @@ class FeatureView:
 
     @property
     def output_schema(self) -> StructType:
-        return self._feature_df.schema
+        return self._infer_schema_df.schema
 
     @property
     def refresh_mode(self) -> Optional[str]:
@@ -285,6 +313,11 @@ class FeatureView:
     def owner(self) -> Optional[str]:
         return self._owner
 
+    def _metadata(self) -> _FeatureViewMetadata:
+        entity_names = [e.name.identifier() for e in self.entities]
+        ts_col = self.timestamp_col.identifier() if self.timestamp_col is not None else _TIMESTAMP_COL_PLACEHOLDER
+        return _FeatureViewMetadata(entity_names, ts_col)
+
     def _get_query(self) -> str:
         if len(self._feature_df.queries["queries"]) != 1:
             raise ValueError(
@@ -300,7 +333,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
                 f"FeatureView name `{self._name}` contains invalid character `{_FEATURE_VIEW_NAME_DELIMITER}`."
             )
 
-        unescaped_df_cols = to_sql_identifiers(self._feature_df.columns)
+        unescaped_df_cols = to_sql_identifiers(self._infer_schema_df.columns)
         for e in self._entities:
             for k in e.join_keys:
                 if k not in unescaped_df_cols:
@@ -312,17 +345,17 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             ts_col = self._timestamp_col
             if ts_col == SqlIdentifier(_TIMESTAMP_COL_PLACEHOLDER):
                 raise ValueError(f"Invalid timestamp_col name, cannot be {_TIMESTAMP_COL_PLACEHOLDER}.")
-            if ts_col not in to_sql_identifiers(self._feature_df.columns):
+            if ts_col not in to_sql_identifiers(self._infer_schema_df.columns):
                 raise ValueError(f"timestamp_col {ts_col} is not found in input dataframe.")
 
-            col_type = self._feature_df.schema[ts_col].datatype
+            col_type = self._infer_schema_df.schema[ts_col].datatype
             if not isinstance(col_type, (DateType, TimeType, TimestampType, _NumericType)):
                 raise ValueError(f"Invalid data type for timestamp_col {ts_col}: {col_type}.")
 
     def _get_feature_names(self) -> List[SqlIdentifier]:
         join_keys = [k for e in self._entities for k in e.join_keys]
         ts_col = [self._timestamp_col] if self._timestamp_col is not None else []
-        feature_names = to_sql_identifiers(self._feature_df.columns, case_sensitive=True)
+        feature_names = to_sql_identifiers(self._infer_schema_df.columns, case_sensitive=False)
         return [c for c in feature_names if c not in join_keys + ts_col]
 
     def __repr__(self) -> str:
@@ -355,6 +388,9 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         fv_dict = self.__dict__.copy()
         if "_feature_df" in fv_dict:
             fv_dict.pop("_feature_df")
+        if "_infer_schema_df" in fv_dict:
+            infer_schema_df = fv_dict.pop("_infer_schema_df")
+            fv_dict["_infer_schema_query"] = infer_schema_df.queries["queries"][0]
         fv_dict["_entities"] = [e._to_dict() for e in self._entities]
         fv_dict["_status"] = str(self._status)
         fv_dict["_name"] = str(self._name) if self._name is not None else None
@@ -411,6 +447,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             refresh_mode=json_dict["_refresh_mode"],
             refresh_mode_reason=json_dict["_refresh_mode_reason"],
             owner=json_dict["_owner"],
+            infer_schema_df=session.sql(json_dict.get("_infer_schema_query", None)),
         )
 
     @staticmethod
@@ -436,12 +473,13 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         status: FeatureViewStatus,
         feature_descs: Dict[str, str],
         refresh_freq: Optional[str],
-        database: Optional[str],
-        schema: Optional[str],
+        database: str,
+        schema: str,
         warehouse: Optional[str],
         refresh_mode: Optional[str],
         refresh_mode_reason: Optional[str],
         owner: Optional[str],
+        infer_schema_df: Optional[DataFrame],
     ) -> FeatureView:
         fv = FeatureView(
             name=name,
@@ -449,6 +487,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             feature_df=feature_df,
             timestamp_col=timestamp_col,
             desc=desc,
+            _infer_schema_df=infer_schema_df,
         )
         fv._version = FeatureViewVersion(version) if version is not None else None
         fv._status = status

snowflake/ml/fileset/embedded_stage_fs.py

@@ -0,0 +1,149 @@
+import re
+from collections import defaultdict
+from typing import Any, List, Optional, Tuple
+
+from snowflake import snowpark
+from snowflake.connector import connection
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.exceptions import (
+    error_codes,
+    exceptions as snowml_exceptions,
+    fileset_errors,
+)
+from snowflake.ml._internal.utils import identifier
+from snowflake.snowpark import exceptions as snowpark_exceptions
+
+from . import stage_fs
+
+_SNOWURL_PATH_RE = re.compile(r"versions/(?P<version>[^/]+)(?:/+(?P<filepath>.*))?")
+
+
+class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
+    def __init__(
+        self,
+        *,
+        domain: str,
+        name: str,
+        snowpark_session: Optional[snowpark.Session] = None,
+        sf_connection: Optional[connection.SnowflakeConnection] = None,
+        **kwargs: Any,
+    ) -> None:
+
+        (db, schema, object_name, _) = identifier.parse_schema_level_object_identifier(name)
+        self._name = name  # TODO: Require or resolve FQN
+        self._domain = domain
+
+        super().__init__(
+            db=db,
+            schema=schema,
+            stage=object_name,
+            snowpark_session=snowpark_session,
+            sf_connection=sf_connection,
+            **kwargs,
+        )
+
+    @property
+    def stage_name(self) -> str:
+        """Get the Snowflake path to this stage.
+
+        Returns:
+            A string in the format of snow://<domain>/<name>
+                Example: snow://dataset/my_dataset
+
+        # noqa: DAR203
+        """
+        return f"snow://{self._domain}/{self._name}"
+
+    def _stage_path_to_relative_path(self, stage_path: str) -> str:
+        """Convert a stage file path which comes from the LIST query to a relative file path in that stage.
+
+        The file path returned by LIST query always has the format "versions/<version>/<relative_file_path>".
+                The full "versions/<version>/<relative_file_path>" is returned
+
+        Args:
+            stage_path: A string started with the name of the stage.
+
+        Returns:
+            A string of the relative stage path.
+        """
+        return stage_path
+
+    def _fetch_presigned_urls(
+        self, files: List[str], url_lifetime: float = stage_fs._PRESIGNED_URL_LIFETIME_SEC
+    ) -> List[Tuple[str, str]]:
+        """Fetch presigned urls for the given files."""
+        # SnowURL requires full snow://<domain>/<entity>/versions/<version> as the stage path arg to get_presigned_Url
+        versions_dict = defaultdict(list)
+        for file in files:
+            match = _SNOWURL_PATH_RE.fullmatch(file)
+            assert match is not None and match.group("filepath") is not None
+            versions_dict[match.group("version")].append(match.group("filepath"))
+        try:
+            async_jobs: List[snowpark.AsyncJob] = []
+            for version, version_files in versions_dict.items():
+                for file in version_files:
+                    stage_loc = f"{self.stage_name}/versions/{version}"
+                    query_result = self._session.sql(
+                        f"select '{version}/{file}' as name,"
+                        f" get_presigned_url('{stage_loc}', '{file}', {url_lifetime}) as url"
+                    ).collect(
+                        block=False,
+                        statement_params=telemetry.get_function_usage_statement_params(
+                            project=stage_fs._PROJECT,
+                            api_calls=[snowpark.DataFrame.collect],
+                        ),
+                    )
+                    async_jobs.append(query_result)
+            presigned_urls: List[Tuple[str, str]] = [
+                (r["NAME"], r["URL"]) for job in async_jobs for r in stage_fs._resolve_async_job(job)
+            ]
+            return presigned_urls
+        except snowpark_exceptions.SnowparkClientException as e:
+            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST) or e.message.startswith(
+                fileset_errors.ERRNO_STAGE_NOT_EXIST
+            ):
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.SNOWML_NOT_FOUND,
+                    original_exception=fileset_errors.StageNotFoundError(
+                        f"Stage {self.stage_name} does not exist or is not authorized."
+                    ),
+                )
+            else:
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.INTERNAL_SNOWML_ERROR,
+                    original_exception=fileset_errors.FileSetError(str(e)),
+                )
+
+    @classmethod
+    def _parent(cls, path: str) -> str:
+        """Get parent of specified path up to minimally valid root path.
+
+        For SnowURL, the minimum valid path is snow://<domain>/<entity>/versions/<version>
+
+        Args:
+            path: File or directory path
+
+        Returns:
+            Parent path
+
+        Examples:
+        ----
+        >>> fs._parent("snow://dataset/my_ds/versions/my_version/file.ext")
+        "snow://dataset/my_ds/versions/my_version/"
+        >>> fs._parent("snow://dataset/my_ds/versions/my_version/subdir/file.ext")
+        "snow://dataset/my_ds/versions/my_version/subdir/"
+        >>> fs._parent("snow://dataset/my_ds/versions/my_version/")
+        "snow://dataset/my_ds/versions/my_version/"
+        >>> fs._parent("snow://dataset/my_ds/versions/my_version")
+        "snow://dataset/my_ds/versions/my_version"
+        """
+        path_match = _SNOWURL_PATH_RE.fullmatch(path)
+        if not path_match:
+            return super()._parent(path)  # type: ignore[no-any-return]
+        filepath: str = path_match.group("filepath") or ""
+        root: str = path[: path_match.start("filepath")] if filepath else path
+        if "/" in filepath:
+            parent = filepath.rsplit("/", 1)[0]
+            return root + parent
+        else:
+            return root

snowflake/ml/fileset/sfcfs.py

@@ -185,7 +185,6 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         func_params_to_log=["detail"],
         conn_attr_name="_conn",
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def ls(self, path: str, detail: bool = False, **kwargs: Any) -> Union[List[str], List[Dict[str, Any]]]:
         """Override fsspec `ls` method. List single "directory" with or without details.
 
@@ -216,7 +215,6 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         project=_PROJECT,
         conn_attr_name="_conn",
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def optimize_read(self, files: Optional[List[str]] = None) -> None:
         """Prefetch and cache the presigned urls for all the given files to speed up the file opening.
 
@@ -242,7 +240,6 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         project=_PROJECT,
         conn_attr_name="_conn",
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def _open(self, path: str, **kwargs: Any) -> fsspec.spec.AbstractBufferedFile:
         """Override fsspec `_open` method. Open a file for reading in 'rb' mode.
 
@@ -268,7 +265,6 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         project=_PROJECT,
         conn_attr_name="_conn",
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def info(self, path: str, **kwargs: Any) -> Dict[str, Any]:
         """Override fsspec `info` method. Give details of entry at path."""
         file_path = self._parse_file_path(path)

snowflake/ml/fileset/snowfs.py

@@ -0,0 +1,160 @@
+import collections
+import logging
+import re
+from typing import Any, Dict, Optional
+
+import fsspec
+import packaging.version as pkg_version
+
+from snowflake import snowpark
+from snowflake.connector import connection
+from snowflake.ml._internal.exceptions import (
+    error_codes,
+    exceptions as snowml_exceptions,
+)
+from snowflake.ml._internal.utils import identifier, snowflake_env
+from snowflake.ml.fileset import embedded_stage_fs, sfcfs
+
+PROTOCOL_NAME = "snow"
+
+_SFFileEntityPath = collections.namedtuple(
+    "_SFFileEntityPath", ["domain", "name", "filepath", "version", "relative_path"]
+)
+_PROJECT = "FileSet"
+_SNOWURL_PATTERN = re.compile(
+    f"({PROTOCOL_NAME}://)?"
+    r"(?<!@)(?P<domain>\w+)/"
+    rf"(?P<name>(?:{identifier._SF_IDENTIFIER}\.){{,2}}{identifier._SF_IDENTIFIER})/"
+    r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)"
+)
+
+# FIXME(dhung): Temporary fix for bug in GS version 8.17
+_BUG_VERSION_MIN = pkg_version.Version("8.17")  # Inclusive minimum version with bugged behavior
+_BUG_VERSION_MAX = pkg_version.Version("8.18")  # Exclusive maximum version with bugged behavior
+
+
+class SnowFileSystem(sfcfs.SFFileSystem):
+    """A filesystem that allows user to access Snowflake embedded stage files with valid Snowflake locations.
+
+    The file system is is based on fsspec (https://filesystem-spec.readthedocs.io/). It is a file system wrapper
+    built on top of SFStageFileSystem. It takes Snowflake embedded stage path as the input and supports read operation.
+    A valid Snowflake location will have the form "snow://{domain}/{entity_name}/versions/{version}/{path_to_file}".
+
+    See `sfcfs.SFFileSystem` documentation for example usage patterns.
+    """
+
+    protocol = PROTOCOL_NAME
+    _IS_BUGGED_VERSION = None
+
+    def __init__(
+        self,
+        sf_connection: Optional[connection.SnowflakeConnection] = None,
+        snowpark_session: Optional[snowpark.Session] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(sf_connection=sf_connection, snowpark_session=snowpark_session, **kwargs)
+
+        # FIXME(dhung): Temporary fix for bug in GS version 8.17
+        if SnowFileSystem._IS_BUGGED_VERSION is None:
+            try:
+                sf_version = snowflake_env.get_current_snowflake_version(self._session)
+                SnowFileSystem._IS_BUGGED_VERSION = _BUG_VERSION_MIN <= sf_version < _BUG_VERSION_MAX
+            except Exception:
+                SnowFileSystem._IS_BUGGED_VERSION = False
+
+    def info(self, path: str, **kwargs: Any) -> Dict[str, Any]:
+        # FIXME(dhung): Temporary fix for bug in GS version 8.17
+        res: Dict[str, Any] = super().info(path, **kwargs)
+        if res.get("type") == "directory" and not res["name"].endswith("/"):
+            res["name"] += "/"
+        return res
+
+    def _get_stage_fs(
+        self, sf_file_path: _SFFileEntityPath  # type: ignore[override]
+    ) -> embedded_stage_fs.SFEmbeddedStageFileSystem:
+        """Get the stage file system for the given snowflake location.
+
+        Args:
+            sf_file_path: The Snowflake path information.
+
+        Returns:
+            A SFEmbeddedStageFileSystem object which supports readonly file operations on Snowflake embedded stages.
+        """
+        stage_fs_key = (sf_file_path.domain, sf_file_path.name, sf_file_path.version)
+        if stage_fs_key not in self._stage_fs_set:
+            cnt_stage_fs = embedded_stage_fs.SFEmbeddedStageFileSystem(
+                snowpark_session=self._session,
+                domain=sf_file_path.domain,
+                name=sf_file_path.name,
+                **self._kwargs,
+            )
+            self._stage_fs_set[stage_fs_key] = cnt_stage_fs
+        return self._stage_fs_set[stage_fs_key]
+
+    def _stage_path_to_absolute_path(self, stage_fs: embedded_stage_fs.SFEmbeddedStageFileSystem, path: str) -> str:
+        """Convert the relative path in a stage to an absolute path starts with the location of the stage."""
+        # Strip protocol from absolute path, since backend needs snow:// prefix to resolve correctly
+        # but fsspec logic strips protocol when doing any searching and globbing
+        stage_name = stage_fs.stage_name
+        protocol = f"{PROTOCOL_NAME}://"
+        if stage_name.startswith(protocol):
+            stage_name = stage_name[len(protocol) :]
+        abs_path = stage_name + "/" + path
+        # FIXME(dhung): Temporary fix for bug in GS version 8.17
+        if self._IS_BUGGED_VERSION:
+            match = _SNOWURL_PATTERN.fullmatch(abs_path)
+            assert match is not None
+            if match.group("relpath"):
+                abs_path = abs_path.replace(match.group("relpath"), match.group("relpath").lstrip("/"))
+        return abs_path
+
+    @classmethod
+    def _parse_file_path(cls, path: str) -> _SFFileEntityPath:  # type: ignore[override]
+        """Parse a snowflake location path.
+
+        The following properties will be extracted from the path input:
+        - embedded stage domain
+        - entity name
+        - path (in format `versions/{version}/{relative_path}`)
+        - entity version (optional)
+        - relative file path (optional)
+
+        Args:
+            path: A string in the format of "snow://{domain}/{entity_name}/versions/{version}/{path_to_file}".
+
+        Returns:
+            A namedtuple consists of domain, entity name, filepath, version, and relative path, where
+                filepath = "versions/{version}/{relative_path}"
+
+        Raises:
+            SnowflakeMLException: An error occurred when invalid path is given.
+        """
+        snowurl_match = _SNOWURL_PATTERN.fullmatch(path)
+        if not snowurl_match:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.SNOWML_INVALID_STAGE,
+                original_exception=ValueError(f"Invalid Snow URL: {path}"),
+            )
+
+        try:
+            domain = snowurl_match.group("domain")
+            parsed_name = identifier.parse_schema_level_object_identifier(snowurl_match.group("name"))
+            name = identifier.get_schema_level_object_identifier(*parsed_name)
+            filepath = snowurl_match.group("path")
+            version = snowurl_match.group("version")
+            relative_path = snowurl_match.group("relpath") or ""
+            logging.debug(f"Parsed snow URL: {snowurl_match.groups()}")
+            # FIXME(dhung): Temporary fix for bug in GS version 8.17
+            if cls._IS_BUGGED_VERSION:
+                filepath = f"versions/{version}//{relative_path}"
+            return _SFFileEntityPath(
+                domain=domain, name=name, version=version, relative_path=relative_path, filepath=filepath
+            )
+        except ValueError as e:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.SNOWML_INVALID_STAGE,
+                original_exception=e,
+            )
+
+
+fsspec.register_implementation(PROTOCOL_NAME, SnowFileSystem)

snowflake/ml/fileset/stage_fs.py

@@ -2,13 +2,13 @@ import inspect
 import logging
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 import fsspec
 from fsspec.implementations import http as httpfs
 
 from snowflake import snowpark
-from snowflake.connector import connection, errorcode
+from snowflake.connector import connection, errorcode, errors as snowpark_errors
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -18,6 +18,7 @@ from snowflake.ml._internal.exceptions import (
 )
 from snowflake.snowpark import exceptions as snowpark_exceptions
 from snowflake.snowpark._internal import utils as snowpark_utils
+from snowflake.snowpark._internal.analyzer import snowflake_plan
 
 # The default length of how long a presigned url stays active in seconds.
 # Presigned url here is used to fetch file objects from Snowflake when SFStageFileSystem.open() is called.
@@ -144,7 +145,6 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         project=_PROJECT,
         func_params_to_log=["detail"],
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def ls(self, path: str, detail: bool = False) -> Union[List[str], List[Dict[str, Any]]]:
         """Override fsspec `ls` method. List single "directory" with or without details.
 
@@ -168,7 +168,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         try:
             loc = self.stage_name
             path = path.lstrip("/")
-            objects = self._session.sql(f"LIST {loc}/{path}").collect()
+            async_job: snowpark.AsyncJob = self._session.sql(f"LIST '{loc}/{path}'").collect(block=False)
+            objects: List[snowpark.Row] = _resolve_async_job(async_job)
         except snowpark_exceptions.SnowparkClientException as e:
             if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST):
                 raise snowml_exceptions.SnowflakeMLException(
@@ -191,7 +192,6 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def optimize_read(self, files: Optional[List[str]] = None) -> None:
         """Prefetch and cache the presigned urls for all the given files to speed up the read performance.
 
@@ -218,7 +218,6 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def _open(self, path: str, mode: str = "rb", **kwargs: Any) -> fsspec.spec.AbstractBufferedFile:
         """Override fsspec `_open` method. Open a file for reading.
 
@@ -292,9 +291,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                     original_exception=e,
                 )
 
-    def _parse_list_result(
-        self, list_result: List[Tuple[str, int, str, str]], search_path: str
-    ) -> List[Dict[str, Any]]:
+    def _parse_list_result(self, list_result: List[snowpark.Row], search_path: str) -> List[Dict[str, Any]]:
         """Convert the result from LIST query to the expected format of fsspec ls() method.
 
         Note that Snowflake LIST query has different behavior with ls(). LIST query will return all the stage files
@@ -315,7 +312,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         """
         files: Dict[str, Dict[str, Any]] = {}
         search_path = search_path.strip("/")
-        for name, size, md5, last_modified in list_result:
+        for row in list_result:
+            name, size, md5, last_modified = row["name"], row["size"], row["md5"], row["last_modified"]
             obj_path = self._stage_path_to_relative_path(name)
             if obj_path == search_path:
                 # If there is a exact match, then the matched object will always be a file object.
@@ -411,3 +409,20 @@ def _match_error_code(ex: snowpark_exceptions.SnowparkSQLException, error_code:
     # Snowpark writes error code to message instead of populating e.error_code
     error_code_str = str(error_code)
     return ex.error_code == error_code_str or error_code_str in ex.message
+
+
+@snowflake_plan.SnowflakePlan.Decorator.wrap_exception  # type: ignore[misc]
+def _resolve_async_job(async_job: snowpark.AsyncJob) -> List[snowpark.Row]:
+    # Make sure Snowpark exceptions are properly caught and converted by wrap_exception wrapper
+    try:
+        query_result = cast(List[snowpark.Row], async_job.result("row"))
+        return query_result
+    except snowpark_errors.DatabaseError as e:
+        # HACK: Snowpark surfaces a generic exception if query doesn't complete immediately
+        # assume it's due to FileNotFound
+        if type(e) is snowpark_errors.DatabaseError and "results are unavailable" in str(e):
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.SNOWML_NOT_FOUND,
+                original_exception=fileset_errors.StageNotFoundError("Query failed."),
+            ) from e
+        raise

snowflake/ml/model/__init__.py

@@ -1,6 +1,6 @@
 from snowflake.ml.model._client.model.model_impl import Model
-from snowflake.ml.model._client.model.model_version_impl import ModelVersion
+from snowflake.ml.model._client.model.model_version_impl import ExportMode, ModelVersion
 from snowflake.ml.model.models.huggingface_pipeline import HuggingFacePipelineModel
 from snowflake.ml.model.models.llm import LLM, LLMOptions
 
-__all__ = ["Model", "ModelVersion", "HuggingFacePipelineModel", "LLM", "LLMOptions"]
+__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "LLM", "LLMOptions"]