snowflake-ml-python 1.5.0__py3-none-any.whl → 1.5.2__py3-none-any.whl

snowflake/ml/feature_store/feature_store.py

@@ -8,7 +8,19 @@ import re
 import warnings
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union, cast
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
 import packaging.version as pkg_version
 import snowflake.ml.version as snowml_version
@@ -32,7 +44,7 @@ from snowflake.ml.feature_store.entity import _ENTITY_NAME_LENGTH_LIMIT, Entity
 from snowflake.ml.feature_store.feature_view import (
     _FEATURE_OBJ_TYPE,
     _FEATURE_VIEW_NAME_DELIMITER,
-    _TIMESTAMP_COL_PLACEHOLDER,
+    _LEGACY_TIMESTAMP_COL_PLACEHOLDER_VALS,
     FeatureView,
     FeatureViewSlice,
     FeatureViewStatus,
@@ -242,23 +254,16 @@ class FeatureStore:
 
         else:
             try:
-                self._session.sql(f"CREATE SCHEMA IF NOT EXISTS {self._config.full_schema_path}").collect(
-                    statement_params=self._telemetry_stmp
-                )
-                for tag in to_sql_identifiers(
-                    [
-                        _FEATURE_VIEW_METADATA_TAG,
-                    ]
-                ):
+                # Explicitly check if schema exists first since we may not have CREATE SCHEMA privilege
+                if len(self._find_object("SCHEMAS", self._config.schema)) == 0:
+                    self._session.sql(f"CREATE SCHEMA IF NOT EXISTS {self._config.full_schema_path}").collect(
+                        statement_params=self._telemetry_stmp
+                    )
+                for tag in to_sql_identifiers([_FEATURE_VIEW_METADATA_TAG, _FEATURE_STORE_OBJECT_TAG]):
                     self._session.sql(f"CREATE TAG IF NOT EXISTS {self._get_fully_qualified_name(tag)}").collect(
                         statement_params=self._telemetry_stmp
                     )
-
-                self._session.sql(
-                    f"CREATE TAG IF NOT EXISTS {self._get_fully_qualified_name(_FEATURE_STORE_OBJECT_TAG)}"
-                ).collect(statement_params=self._telemetry_stmp)
             except Exception as e:
-                self.clear()
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
                     original_exception=RuntimeError(f"Failed to create feature store {name}: {e}."),
@@ -750,7 +755,7 @@ class FeatureStore:
         except Exception as e:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
-                original_exception=RuntimeError(f"Failed to alter schema or drop tag: {e}."),
+                original_exception=RuntimeError(f"Failed to delete entity: {e}."),
             ) from e
         logger.info(f"Deleted Entity {name}.")
 
@@ -802,7 +807,7 @@ class FeatureStore:
 
         return df
 
-    @dispatch_decorator()
+    @overload
     def generate_dataset(
         self,
         name: str,
@@ -814,7 +819,40 @@ class FeatureStore:
         exclude_columns: Optional[List[str]] = None,
         include_feature_view_timestamp_col: bool = False,
         desc: str = "",
+        output_type: Literal["dataset"] = "dataset",
     ) -> dataset.Dataset:
+        ...
+
+    @overload
+    def generate_dataset(
+        self,
+        name: str,
+        spine_df: DataFrame,
+        features: List[Union[FeatureView, FeatureViewSlice]],
+        output_type: Literal["table"],
+        version: Optional[str] = None,
+        spine_timestamp_col: Optional[str] = None,
+        spine_label_cols: Optional[List[str]] = None,
+        exclude_columns: Optional[List[str]] = None,
+        include_feature_view_timestamp_col: bool = False,
+        desc: str = "",
+    ) -> DataFrame:
+        ...
+
+    @dispatch_decorator()  # type: ignore[misc]
+    def generate_dataset(
+        self,
+        name: str,
+        spine_df: DataFrame,
+        features: List[Union[FeatureView, FeatureViewSlice]],
+        version: Optional[str] = None,
+        spine_timestamp_col: Optional[str] = None,
+        spine_label_cols: Optional[List[str]] = None,
+        exclude_columns: Optional[List[str]] = None,
+        include_feature_view_timestamp_col: bool = False,
+        desc: str = "",
+        output_type: Literal["dataset", "table"] = "dataset",
+    ) -> Union[dataset.Dataset, DataFrame]:
         """
         Generate dataset by given source table and feature views.
 
@@ -834,30 +872,29 @@ class FeatureStore:
             include_feature_view_timestamp_col: Generated dataset will include timestamp column of feature view
                 (if feature view has timestamp column) if set true. Default to false.
             desc: A description about this dataset.
+            output_type: The type of Snowflake storage to use for the generated training data.
 
         Returns:
-            A Dataset object.
+            If output_type is "dataset" (default), returns a Dataset object.
+            If output_type is "table", returns a Snowpark DataFrame representing the table.
 
         Raises:
-            SnowflakeMLException: [ValueError] spine_df contains more than one query.
             SnowflakeMLException: [ValueError] Dataset name/version already exists
             SnowflakeMLException: [ValueError] Snapshot creation failed.
+            SnowflakeMLException: [ValueError] Invalid output_type specified.
             SnowflakeMLException: [RuntimeError] Failed to create clone from table.
             SnowflakeMLException: [RuntimeError] Failed to find resources.
         """
+        if output_type not in {"table", "dataset"}:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_ARGUMENT,
+                original_exception=ValueError(f"Invalid output_type: {output_type}."),
+            )
         if spine_timestamp_col is not None:
             spine_timestamp_col = SqlIdentifier(spine_timestamp_col)
         if spine_label_cols is not None:
             spine_label_cols = to_sql_identifiers(spine_label_cols)  # type: ignore[assignment]
 
-        if len(spine_df.queries["queries"]) != 1:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    f"spine_df must contain only one query. Got: {spine_df.queries['queries']}"
-                ),
-            )
-
         result_df, join_keys = self._join_features(
             spine_df, features, spine_timestamp_col, include_feature_view_timestamp_col
         )
@@ -875,33 +912,49 @@ class FeatureStore:
             result_df = self._exclude_columns(result_df, exclude_columns)
 
         fs_meta = FeatureStoreMetadata(
-            spine_query=spine_df.queries["queries"][0],
+            spine_query=spine_df.queries["queries"][-1],
             serialized_feature_views=[fv.to_json() for fv in features],
             spine_timestamp_col=spine_timestamp_col,
         )
 
         try:
-            ds: dataset.Dataset = dataset.create_from_dataframe(
-                self._session,
-                name,
-                version,
-                input_dataframe=result_df,
-                exclude_cols=[spine_timestamp_col],
-                label_cols=spine_label_cols,
-                properties=fs_meta,
-                comment=desc,
-            )
-            return ds
+            if output_type == "table":
+                table_name = f"{name}_{version}"
+                result_df.write.mode("errorifexists").save_as_table(table_name)  # type: ignore[call-overload]
+                ds_df = self._session.table(table_name)
+                return ds_df
+            else:
+                assert output_type == "dataset"
+                if not self._is_dataset_enabled():
+                    raise snowml_exceptions.SnowflakeMLException(
+                        error_code=error_codes.SNOWML_CREATE_FAILED,
+                        original_exception=RuntimeError(
+                            "Dataset is not enabled in your account. Ask your account admin to set"
+                            ' FEATURE_DATASET=ENABLED or set output_type="table" to generate the data'
+                            " as a Snowflake Table instead."
+                        ),
+                    )
+                ds: dataset.Dataset = dataset.create_from_dataframe(
+                    self._session,
+                    name,
+                    version,
+                    input_dataframe=result_df,
+                    exclude_cols=[spine_timestamp_col],
+                    label_cols=spine_label_cols,
+                    properties=fs_meta,
+                    comment=desc,
+                )
+                return ds
 
         except dataset_errors.DatasetExistError as e:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.OBJECT_ALREADY_EXISTS,
-                original_exception=ValueError(str(e)),
+                original_exception=RuntimeError(str(e)),
             ) from e
         except SnowparkSQLException as e:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
-                original_exception=RuntimeError(f"An error occurred during Dataset generation: {e}."),
+                original_exception=RuntimeError(f"An error occurred during dataset generation: {e}."),
             ) from e
 
     @dispatch_decorator()
@@ -930,52 +983,47 @@ class FeatureStore:
         return self._load_serialized_feature_objects(source_meta.properties.serialized_feature_views)
 
     @dispatch_decorator()
-    def clear(self) -> None:
+    def _clear(self, dryrun: bool = True) -> None:
         """
-        Clear all feature store internal objects including feature views, entities etc. Note feature store
-        instance (snowflake schema) won't be deleted. Use snowflake to delete feature store instance.
+        Clear all feature views and entities. Note Feature Store schema and metadata will NOT be purged
+        together. Use SQL to delete schema and metadata instead.
 
-        Raises:
-            SnowflakeMLException: [RuntimeError] Failed to clear feature store.
+        Args:
+            dryrun: Print a list of objects will be deleted but not actually perform the deletion when true.
         """
-        try:
-            result = self._session.sql(
-                f"""
-                SELECT *
-                FROM {self._config.database}.INFORMATION_SCHEMA.SCHEMATA
-                WHERE SCHEMA_NAME = '{self._config.schema.resolved()}'
-            """
-            ).collect()
-            if len(result) == 0:
-                return
-
-            fs_obj_tag = self._find_object("TAGS", SqlIdentifier(_FEATURE_STORE_OBJECT_TAG))
-            if len(fs_obj_tag) == 0:
-                return
-
-            object_types = ["DYNAMIC TABLES", "DATASETS", "VIEWS", "TASKS"]
-            for obj_type in object_types:
-                all_object_rows = self._find_object(obj_type, None)
-                for row in all_object_rows:
-                    obj_name = self._get_fully_qualified_name(SqlIdentifier(row["name"], case_sensitive=True))
-                    self._session.sql(f"DROP {obj_type[:-1]} {obj_name}").collect()
-                    logger.info(f"Deleted {obj_type[:-1]}: {obj_name}.")
-
-            entity_tags = self._find_object("TAGS", SqlIdentifier(_ENTITY_TAG_PREFIX), prefix_match=True)
-            all_tags = [
-                _FEATURE_STORE_OBJECT_TAG,
-                _FEATURE_VIEW_METADATA_TAG,
-            ] + [SqlIdentifier(row["name"], case_sensitive=True) for row in entity_tags]
-            for tag_name in all_tags:
-                obj_name = self._get_fully_qualified_name(tag_name)
-                self._session.sql(f"DROP TAG IF EXISTS {obj_name}").collect()
-                logger.info(f"Deleted TAG: {obj_name}.")
+        warnings.warn(
+            "It will clear ALL feature views and entities in this Feature Store. Make sure your role"
+            " has sufficient access to all feature views and entities. Insufficient access to some feature"
+            " views or entities will leave Feature Store in an incomplete state.",
+            stacklevel=2,
+            category=UserWarning,
+        )
+
+        all_fvs_df = self.list_feature_views()
+        all_entities_df = self.list_entities()
+        all_fvs_rows = all_fvs_df.collect()
+        all_entities_rows = all_entities_df.collect()
+
+        if dryrun:
+            logger.info(
+                "Following feature views and entities will be deleted."
+                + " Set 'dryrun=False' to perform the actual deletion."
+            )
+            logger.info(f"Total {len(all_fvs_rows)} Feature views to be deleted:")
+            all_fvs_df.show(n=len(all_fvs_rows))
+            logger.info(f"\nTotal {len(all_entities_rows)} entities to be deleted:")
+            all_entities_df.show(n=len(all_entities_rows))
+            return
+
+        for fv_row in all_fvs_rows:
+            fv = self.get_feature_view(
+                SqlIdentifier(fv_row["NAME"], case_sensitive=True).identifier(), fv_row["VERSION"]
+            )
+            self.delete_feature_view(fv)
+
+        for entity_row in all_entities_rows:
+            self.delete_entity(SqlIdentifier(entity_row["NAME"], case_sensitive=True).identifier())
 
-        except Exception as e:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
-                original_exception=RuntimeError(f"Failed to clear feature store {self._config.full_schema_path}: {e}."),
-            ) from e
         logger.info(f"Feature store {self._config.full_schema_path} has been cleared.")
 
     def _get_feature_view_if_exists(self, name: str, version: str) -> FeatureView:
@@ -1093,14 +1141,6 @@ class FeatureStore:
         spine_timestamp_col: Optional[SqlIdentifier],
         include_feature_view_timestamp_col: bool,
     ) -> Tuple[DataFrame, List[SqlIdentifier]]:
-        if len(spine_df.queries["queries"]) != 1:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    f"spine_df must contain only one query. Got: {spine_df.queries['queries']}"
-                ),
-            )
-
         for f in features:
             f = f.feature_view_ref if isinstance(f, FeatureViewSlice) else f
             if f.status == FeatureViewStatus.DRAFT:
@@ -1122,7 +1162,7 @@ class FeatureStore:
             self._asof_join_enabled = self._is_asof_join_enabled()
 
         # TODO: leverage Snowpark dataframe for more concise syntax once it supports AsOfJoin
-        query = spine_df.queries["queries"][0]
+        query = spine_df.queries["queries"][-1]
         layer = 0
         for f in features:
             if isinstance(f, FeatureViewSlice):
@@ -1180,7 +1220,15 @@ class FeatureStore:
                 """
             layer += 1
 
-        return self._session.sql(query), join_keys
+        # TODO: construct result dataframe with datframe APIs once ASOF join is supported natively.
+        # Below code manually construct result dataframe from private members of spine dataframe, which
+        # likely will cause unintentional issues. This setp is needed because spine_df might contains
+        # prerequisite queries and post actions that must be carried over to result dataframe.
+        result_df = self._session.sql(query)
+        result_df._plan.queries = spine_df._plan.queries[:-1] + result_df._plan.queries
+        result_df._plan.post_actions = spine_df._plan.post_actions
+
+        return result_df, join_keys
 
     def _check_database_exists_or_throw(self) -> None:
         resolved_db_name = self._config.database.resolved()
@@ -1517,6 +1565,9 @@ class FeatureStore:
                 original_exception=RuntimeError(f"Failed to parse query text for FeatureView {name}/{version}: {row}."),
             )
 
+        fv_name = FeatureView._get_physical_name(name, version)
+        infer_schema_df = self._session.sql(f"SELECT * FROM {self._get_fully_qualified_name(fv_name)}")
+
         if m.group("obj_type") == "DYNAMIC TABLE":
             query = m.group("query")
             df = self._session.sql(query)
@@ -1524,7 +1575,7 @@ class FeatureStore:
             fv_metadata = _FeatureViewMetadata.from_json(m.group("fv_metadata"))
             entities = [find_and_compose_entity(n) for n in fv_metadata.entities]
             ts_col = fv_metadata.timestamp_col
-            timestamp_col = ts_col if ts_col != _TIMESTAMP_COL_PLACEHOLDER else None
+            timestamp_col = ts_col if ts_col not in _LEGACY_TIMESTAMP_COL_PLACEHOLDER_VALS else None
 
             fv = FeatureView._construct_feature_view(
                 name=name,
@@ -1534,9 +1585,7 @@ class FeatureStore:
                 desc=desc,
                 version=version,
                 status=FeatureViewStatus(row["scheduling_state"]),
-                feature_descs=self._fetch_column_descs(
-                    "DYNAMIC TABLE", SqlIdentifier(row["name"], case_sensitive=True)
-                ),
+                feature_descs=self._fetch_column_descs("DYNAMIC TABLE", fv_name),
                 refresh_freq=row["target_lag"],
                 database=self._config.database.identifier(),
                 schema=self._config.schema.identifier(),
@@ -1544,6 +1593,7 @@ class FeatureStore:
                 refresh_mode=row["refresh_mode"],
                 refresh_mode_reason=row["refresh_mode_reason"],
                 owner=row["owner"],
+                infer_schema_df=infer_schema_df,
             )
             return fv
         else:
@@ -1553,7 +1603,7 @@ class FeatureStore:
             fv_metadata = _FeatureViewMetadata.from_json(m.group("fv_metadata"))
             entities = [find_and_compose_entity(n) for n in fv_metadata.entities]
             ts_col = fv_metadata.timestamp_col
-            timestamp_col = ts_col if ts_col != _TIMESTAMP_COL_PLACEHOLDER else None
+            timestamp_col = ts_col if ts_col not in _LEGACY_TIMESTAMP_COL_PLACEHOLDER_VALS else None
 
             fv = FeatureView._construct_feature_view(
                 name=name,
@@ -1563,7 +1613,7 @@ class FeatureStore:
                 desc=desc,
                 version=version,
                 status=FeatureViewStatus.STATIC,
-                feature_descs=self._fetch_column_descs("VIEW", SqlIdentifier(row["name"], case_sensitive=True)),
+                feature_descs=self._fetch_column_descs("VIEW", fv_name),
                 refresh_freq=None,
                 database=self._config.database.identifier(),
                 schema=self._config.schema.identifier(),
@@ -1571,6 +1621,7 @@ class FeatureStore:
                 refresh_mode=None,
                 refresh_mode_reason=None,
                 owner=row["owner"],
+                infer_schema_df=infer_schema_df,
             )
             return fv
 
@@ -1710,7 +1761,7 @@ class FeatureStore:
             self._session.sql(
                 f"""
                     SELECT * FROM TABLE(
-                        INFORMATION_SCHEMA.TAG_REFERENCES_INTERNAL(
+                        {self._config.database}.INFORMATION_SCHEMA.TAG_REFERENCES_INTERNAL(
                             TAG_NAME => '{_FEATURE_STORE_OBJECT_TAG}'
                         )
                     ) LIMIT 1;
@@ -1720,6 +1771,15 @@ class FeatureStore:
         except Exception:
             return False
 
+    def _is_dataset_enabled(self) -> bool:
+        try:
+            self._session.sql(f"SHOW DATASETS IN SCHEMA {self._config.full_schema_path}").collect()
+            return True
+        except SnowparkSQLException as e:
+            if "'DATASETS' does not exist" in e.message:
+                return False
+            raise
+
     def _check_feature_store_object_versions(self) -> None:
         versions = self._collapse_object_versions()
         if len(versions) > 0 and pkg_version.parse(snowml_version.VERSION) < versions[0]:

snowflake/ml/feature_store/feature_view.py

@@ -5,12 +5,8 @@ import re
 from collections import OrderedDict
 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,
-    exceptions as snowml_exceptions,
-)
 from snowflake.ml._internal.utils.identifier import concat_names
 from snowflake.ml._internal.utils.sql_identifier import (
     SqlIdentifier,
@@ -27,12 +23,18 @@ 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
 
+_RESULT_SCAN_QUERY_PATTERN = re.compile(
+    r".*FROM\s*TABLE\s*\(\s*RESULT_SCAN\s*\(.*",
+    flags=re.DOTALL | re.IGNORECASE | re.X,
+)
+
 
 @dataclass(frozen=True)
 class _FeatureViewMetadata:
@@ -53,13 +55,10 @@ class _FeatureViewMetadata:
 class FeatureViewVersion(str):
     def __new__(cls, version: str) -> FeatureViewVersion:
         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"`{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}."
-                ),
+            raise 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)
 
@@ -121,12 +120,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).
@@ -140,6 +140,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)
@@ -149,6 +150,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
@@ -295,7 +297,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]:
@@ -329,7 +331,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:
@@ -341,17 +343,20 @@ 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}.")
 
+        if re.match(_RESULT_SCAN_QUERY_PATTERN, self._query) is not None:
+            raise ValueError(f"feature_df should not be reading from RESULT_SCAN. Invalid query: {self._query}")
+
     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:
@@ -384,6 +389,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
@@ -440,6 +448,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
@@ -471,6 +480,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         refresh_mode: Optional[str],
         refresh_mode_reason: Optional[str],
         owner: Optional[str],
+        infer_schema_df: Optional[DataFrame],
     ) -> FeatureView:
         fv = FeatureView(
             name=name,
@@ -478,6 +488,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

@@ -78,22 +78,26 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
             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"))
-        presigned_urls: List[Tuple[str, str]] = []
         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}"
-                    presigned_urls.extend(
-                        self._session.sql(
-                            f"select '{version}/{file}' as name,"
-                            f" get_presigned_url('{stage_loc}', '{file}', {url_lifetime}) as url"
-                        ).collect(
-                            statement_params=telemetry.get_function_usage_statement_params(
-                                project=stage_fs._PROJECT,
-                                api_calls=[snowpark.DataFrame.collect],
-                            ),
-                        )
+                    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
@@ -109,7 +113,6 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
                     error_code=error_codes.INTERNAL_SNOWML_ERROR,
                     original_exception=fileset_errors.FileSetError(str(e)),
                 )
-        return presigned_urls
 
     @classmethod
     def _parent(cls, path: str) -> str: