snowflake-ml-python 1.6.0__py3-none-any.whl → 1.6.1__py3-none-any.whl

snowflake/ml/feature_store/feature_store.py

@@ -52,6 +52,7 @@ from snowflake.ml.feature_store.feature_view import (
     FeatureViewVersion,
     _FeatureViewMetadata,
 )
+from snowflake.ml.utils import sql_client
 from snowflake.snowpark import DataFrame, Row, Session, functions as F
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.types import (
@@ -94,13 +95,12 @@ class _FeatureStoreObjInfo:
         return cls(**state_dict)  # type: ignore[arg-type]
 
 
-# TODO: remove "" after dataset is updated
 class _FeatureStoreObjTypes(Enum):
     UNKNOWN = "UNKNOWN"  # for forward compatibility
     MANAGED_FEATURE_VIEW = "MANAGED_FEATURE_VIEW"
     EXTERNAL_FEATURE_VIEW = "EXTERNAL_FEATURE_VIEW"
     FEATURE_VIEW_REFRESH_TASK = "FEATURE_VIEW_REFRESH_TASK"
-    TRAINING_DATA = ""
+    TRAINING_DATA = "TRAINING_DATA"
 
     @classmethod
     def parse(cls, val: str) -> _FeatureStoreObjTypes:
@@ -140,9 +140,8 @@ _LIST_FEATURE_VIEW_SCHEMA = StructType(
 )
 
 
-class CreationMode(Enum):
-    FAIL_IF_NOT_EXIST = 1
-    CREATE_IF_NOT_EXIST = 2
+CreationMode = sql_client.CreationOption
+CreationMode.__module__ = __name__
 
 
 @dataclass(frozen=True)
@@ -426,7 +425,9 @@ class FeatureStore:
 
         """
         name = SqlIdentifier(name)
-        found_rows = self.list_entities().filter(F.col("NAME") == name.resolved()).collect()
+        found_rows = (
+            self.list_entities().filter(F.col("NAME") == name.resolved()).collect(statement_params=self._telemetry_stmp)
+        )
 
         if len(found_rows) == 0:
             warnings.warn(
@@ -853,7 +854,9 @@ class FeatureStore:
                 original_exception=ValueError(f"Failed to find FeatureView {name}/{version}: {results}"),
             )
 
-        return self._compose_feature_view(results[0][0], results[0][1], self.list_entities().collect())
+        return self._compose_feature_view(
+            results[0][0], results[0][1], self.list_entities().collect(statement_params=self._telemetry_stmp)
+        )
 
     @overload
     def refresh_feature_view(self, feature_view: FeatureView) -> None:
@@ -1223,7 +1226,11 @@ class FeatureStore:
         """
         name = SqlIdentifier(name)
         try:
-            result = self.list_entities().filter(F.col("NAME") == name.resolved()).collect()
+            result = (
+                self.list_entities()
+                .filter(F.col("NAME") == name.resolved())
+                .collect(statement_params=self._telemetry_stmp)
+            )
         except Exception as e:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
@@ -1357,7 +1364,7 @@ class FeatureStore:
         if len(features) == 0:
             raise ValueError("features cannot be empty")
         if isinstance(features[0], str):
-            features = self._load_serialized_feature_objects(cast(List[str], features))
+            features = self._load_serialized_feature_views(cast(List[str], features))
 
         df, _ = self._join_features(
             spine_df,
@@ -1441,8 +1448,19 @@ class FeatureStore:
         if save_as is not None:
             try:
                 save_as = self._get_fully_qualified_name(save_as)
-                result_df.write.mode("errorifexists").save_as_table(save_as)
+                result_df.write.mode("errorifexists").save_as_table(save_as, statement_params=self._telemetry_stmp)
+
+                # Add tag
+                task_obj_info = _FeatureStoreObjInfo(_FeatureStoreObjTypes.TRAINING_DATA, snowml_version.VERSION)
+                self._session.sql(
+                    f"""
+                    ALTER TABLE {save_as}
+                    SET TAG {self._get_fully_qualified_name(_FEATURE_STORE_OBJECT_TAG)}='{task_obj_info.to_json()}'
+                    """
+                ).collect(statement_params=self._telemetry_stmp)
+
                 return self._session.table(save_as)
+
             except SnowparkSQLException as e:
                 if e.sql_error_code == sql_error_codes.OBJECT_ALREADY_EXISTS:
                     raise snowml_exceptions.SnowflakeMLException(
@@ -1572,7 +1590,7 @@ class FeatureStore:
 
         fs_meta = FeatureStoreMetadata(
             spine_query=spine_df.queries["queries"][-1],
-            serialized_feature_views=[fv.to_json() for fv in features],
+            compact_feature_views=[fv._get_compact_repr().to_json() for fv in features],
             spine_timestamp_col=spine_timestamp_col,
         )
 
@@ -1607,6 +1625,7 @@ class FeatureStore:
                         " to generate the data as a Snowflake Table."
                     ),
                 )
+            # TODO: Add feature store tag once Dataset (version) supports tags
             ds: dataset.Dataset = dataset.create_from_dataframe(
                 self._session,
                 name,
@@ -1675,11 +1694,18 @@ class FeatureStore:
         if (
             source_meta is None
             or not isinstance(source_meta.properties, FeatureStoreMetadata)
-            or source_meta.properties.serialized_feature_views is None
+            or (
+                source_meta.properties.serialized_feature_views is None
+                and source_meta.properties.compact_feature_views is None
+            )
         ):
             raise ValueError(f"Dataset {ds} does not contain valid feature view information.")
 
-        return self._load_serialized_feature_objects(source_meta.properties.serialized_feature_views)
+        properties = source_meta.properties
+        if properties.serialized_feature_views:
+            return self._load_serialized_feature_views(properties.serialized_feature_views)
+        else:
+            return self._load_compact_feature_views(properties.compact_feature_views)  # type: ignore[arg-type]
 
     @dispatch_decorator()
     def _clear(self, dryrun: bool = True) -> None:
@@ -1700,8 +1726,8 @@ class FeatureStore:
 
         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()
+        all_fvs_rows = all_fvs_df.collect(statement_params=self._telemetry_stmp)
+        all_entities_rows = all_entities_df.collect(statement_params=self._telemetry_stmp)
 
         if dryrun:
             logger.info(
@@ -1768,6 +1794,7 @@ class FeatureStore:
                     {tagging_clause}
                 )
                 WAREHOUSE = {warehouse}
+                REFRESH_MODE = {feature_view.refresh_mode}
                 AS {feature_view.query}
             """
             self._session.sql(query).collect(block=block, statement_params=self._telemetry_stmp)
@@ -1985,7 +2012,7 @@ class FeatureStore:
                 MATCH_CONDITION ( spine.ts >= feature.ts )
                 ON spine.id = feature.id;
             """
-            ).collect()
+            ).collect(statement_params=self._telemetry_stmp)
         except SnowparkSQLException:
             return False
         return result is not None and len(result) == 1
@@ -2366,11 +2393,11 @@ class FeatureStore:
                 result.append(row)
         return result
 
-    def _load_serialized_feature_objects(
-        self, serialized_feature_objs: List[str]
+    def _load_serialized_feature_views(
+        self, serialized_feature_views: List[str]
     ) -> List[Union[FeatureView, FeatureViewSlice]]:
         results: List[Union[FeatureView, FeatureViewSlice]] = []
-        for obj in serialized_feature_objs:
+        for obj in serialized_feature_views:
             try:
                 obj_type = json.loads(obj)[_FEATURE_OBJ_TYPE]
             except Exception as e:
@@ -2384,6 +2411,14 @@ class FeatureStore:
                 raise ValueError(f"Unsupported feature object type: {obj_type}")
         return results
 
+    def _load_compact_feature_views(
+        self, compact_feature_views: List[str]
+    ) -> List[Union[FeatureView, FeatureViewSlice]]:
+        results: List[Union[FeatureView, FeatureViewSlice]] = []
+        for obj in compact_feature_views:
+            results.append(FeatureView._load_from_compact_repr(self._session, obj))
+        return results
+
     def _exclude_columns(self, df: DataFrame, exclude_columns: List[str]) -> DataFrame:
         exclude_columns = to_sql_identifiers(exclude_columns)  # type: ignore[assignment]
         df_cols = to_sql_identifiers(df.columns)
@@ -2399,12 +2434,12 @@ class FeatureStore:
 
     def _is_dataset_enabled(self) -> bool:
         try:
-            self._session.sql(f"SHOW DATASETS IN SCHEMA {self._config.full_schema_path}").collect()
+            self._session.sql(f"SHOW DATASETS IN SCHEMA {self._config.full_schema_path}").collect(
+                statement_params=self._telemetry_stmp
+            )
             return True
-        except SnowparkSQLException as e:
-            if "'DATASETS' does not exist" in e.message:
-                return False
-            raise
+        except SnowparkSQLException:
+            return False
 
     def _check_feature_store_object_versions(self) -> None:
         versions = self._collapse_object_versions()

snowflake/ml/feature_store/feature_view.py

@@ -6,7 +6,7 @@ import warnings
 from collections import OrderedDict
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -60,6 +60,29 @@ class _FeatureViewMetadata:
         return cls(**state_dict)
 
 
+@dataclass(frozen=True)
+class _CompactRepresentation:
+    """
+    A compact representation for FeatureView and FeatureViewSlice, which contains fully qualified name
+    and optionally a list of feature indices (None means all features will be included).
+    This is to make the metadata much smaller when generating dataset.
+    """
+
+    db: str
+    sch: str
+    name: str
+    version: str
+    feature_indices: Optional[List[int]] = None
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self))
+
+    @classmethod
+    def from_json(cls, json_str: str) -> _CompactRepresentation:
+        state_dict = json.loads(json_str)
+        return cls(**state_dict)
+
+
 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:
@@ -115,6 +138,19 @@ class FeatureViewSlice:
         json_dict["feature_view_ref"] = FeatureView.from_json(json_dict["feature_view_ref"], session)
         return cls(**json_dict)
 
+    def _get_compact_repr(self) -> _CompactRepresentation:
+        return _CompactRepresentation(
+            db=self.feature_view_ref.database.identifier(),  # type: ignore[union-attr]
+            sch=self.feature_view_ref.schema.identifier(),  # type: ignore[union-attr]
+            name=self.feature_view_ref.name.identifier(),
+            version=self.feature_view_ref.version,  # type: ignore[arg-type]
+            feature_indices=self._feature_names_to_indices(),
+        )
+
+    def _feature_names_to_indices(self) -> List[int]:
+        name_to_indices_map = {name: idx for idx, name in enumerate(self.feature_view_ref.feature_names)}
+        return [name_to_indices_map[n] for n in self.names]
+
 
 class FeatureView(lineage_node.LineageNode):
     """
@@ -196,7 +232,7 @@ class FeatureView(lineage_node.LineageNode):
         self._database: Optional[SqlIdentifier] = None
         self._schema: Optional[SqlIdentifier] = None
         self._warehouse: Optional[SqlIdentifier] = SqlIdentifier(warehouse) if warehouse is not None else None
-        self._refresh_mode: Optional[str] = None
+        self._refresh_mode: Optional[str] = _kwargs.get("refresh_mode", "AUTO")
         self._refresh_mode_reason: Optional[str] = None
         self._owner: Optional[str] = None
         self._validate()
@@ -394,6 +430,54 @@ class FeatureView(lineage_node.LineageNode):
     def feature_descs(self) -> Dict[SqlIdentifier, str]:
         return self._feature_desc
 
+    def list_columns(self) -> DataFrame:
+        """List all columns and their information.
+
+        Returns:
+            A Snowpark DataFrame contains feature information.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = Entity("foo", ["id"], desc='my entity')
+            >>> fs.register_entity(e)
+            <BLANKLINE>
+            >>> draft_fv = FeatureView(
+            ...     name="fv",
+            ...     entities=[e],
+            ...     feature_df=self._session.table(<source_table>).select(["NAME", "ID", "TITLE", "AGE", "TS"]),
+            ...     timestamp_col="ts",
+            >>> ).attach_feature_desc({"AGE": "my age", "TITLE": '"my title"'})
+            >>> fv = fs.register_feature_view(draft_fv, '1.0')
+            <BLANKLINE>
+            >>> fv.list_columns().show()
+            --------------------------------------------------
+            |"NAME"  |"CATEGORY"  |"DTYPE"      |"DESC"      |
+            --------------------------------------------------
+            |NAME    |FEATURE     |string(64)   |            |
+            |ID      |ENTITY      |bigint       |my entity   |
+            |TITLE   |FEATURE     |string(128)  |"my title"  |
+            |AGE     |FEATURE     |bigint       |my age      |
+            |TS      |TIMESTAMP   |bigint       |NULL        |
+            --------------------------------------------------
+
+        """
+        session = self._feature_df.session
+        rows = []
+        for name, type in self._feature_df.dtypes:
+            if SqlIdentifier(name) in self.feature_descs:
+                desc = self.feature_descs[SqlIdentifier(name)]
+                rows.append((name, "FEATURE", type, desc))
+            elif SqlIdentifier(name) == self._timestamp_col:
+                rows.append((name, "TIMESTAMP", type, None))  # type: ignore[arg-type]
+            else:
+                for e in self._entities:
+                    if SqlIdentifier(name) in e.join_keys:
+                        rows.append((name, "ENTITY", type, e.desc))
+                        break
+
+        return session.create_dataframe(rows, schema=["name", "category", "dtype", "desc"])
+
     @property
     def refresh_freq(self) -> Optional[str]:
         return self._refresh_freq
@@ -599,12 +683,50 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
 
         return fv_dict
 
-    def to_df(self, session: Session) -> DataFrame:
+    def to_df(self, session: Optional[Session] = None) -> DataFrame:
+        """Convert feature view to a Snowpark DataFrame object.
+
+        Args:
+            session: [deprecated] This argument has no effect. No need to pass a session object.
+
+        Returns:
+            A Snowpark Dataframe object contains the information about feature view.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e = Entity("foo", ["id"], desc='my entity')
+            >>> fs.register_entity(e)
+            <BLANKLINE>
+            >>> draft_fv = FeatureView(
+            ...     name="fv",
+            ...     entities=[e],
+            ...     feature_df=self._session.table(<source_table>).select(["NAME", "ID", "TITLE", "AGE", "TS"]),
+            ...     timestamp_col="ts",
+            >>> ).attach_feature_desc({"AGE": "my age", "TITLE": '"my title"'})
+            >>> fv = fs.register_feature_view(draft_fv, '1.0')
+            <BLANKLINE>
+            fv.to_df().show()
+            ----------------------------------------------------------------...
+            |"NAME"  |"ENTITIES"                |"TIMESTAMP_COL"  |"DESC"  |
+            ----------------------------------------------------------------...
+            |FV      |[                         |TS               |foobar  |
+            |        |  {                       |                 |        |
+            |        |    "desc": "my entity",  |                 |        |
+            |        |    "join_keys": [        |                 |        |
+            |        |      "ID"                |                 |        |
+            |        |    ],                    |                 |        |
+            |        |    "name": "FOO",        |                 |        |
+            |        |    "owner": null         |                 |        |
+            |        |  }                       |                 |        |
+            |        |]                         |                 |        |
+            ----------------------------------------------------------------...
+        """
         values = list(self._to_dict().values())
         schema = [x.lstrip("_") for x in list(self._to_dict().keys())]
         values.append(str(FeatureView._get_physical_name(self._name, self._version)))  # type: ignore[arg-type]
         schema.append("physical_name")
-        return session.create_dataframe([values], schema=schema)
+        return self._feature_df.session.create_dataframe([values], schema=schema)
 
     def to_json(self) -> str:
         state_dict = self._to_dict()
@@ -643,6 +765,14 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             session=session,
         )
 
+    def _get_compact_repr(self) -> _CompactRepresentation:
+        return _CompactRepresentation(
+            db=self.database.identifier(),  # type: ignore[union-attr]
+            sch=self.schema.identifier(),  # type: ignore[union-attr]
+            name=self.name.identifier(),
+            version=self.version,  # type: ignore[arg-type]
+        )
+
     @staticmethod
     def _get_physical_name(fv_name: SqlIdentifier, fv_version: FeatureViewVersion) -> SqlIdentifier:
         return SqlIdentifier(
@@ -655,6 +785,20 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             )
         )
 
+    @staticmethod
+    def _load_from_compact_repr(session: Session, serialized_repr: str) -> Union[FeatureView, FeatureViewSlice]:
+        compact_repr = _CompactRepresentation.from_json(serialized_repr)
+
+        fs = feature_store.FeatureStore(
+            session, compact_repr.db, compact_repr.sch, default_warehouse=session.get_current_warehouse()
+        )
+        fv = fs.get_feature_view(compact_repr.name, compact_repr.version)
+
+        if compact_repr.feature_indices is not None:
+            feature_names = [fv.feature_names[i] for i in compact_repr.feature_indices]
+            return fv.slice(feature_names)  # type: ignore[no-any-return]
+        return fv  # type: ignore[no-any-return]
+
     @staticmethod
     def _load_from_lineage_node(session: Session, name: str, version: str) -> FeatureView:
         db_name, feature_store_name, feature_view_name, _ = identifier.parse_schema_level_object_identifier(name)

snowflake/ml/model/_client/model/model_impl.py

@@ -5,7 +5,7 @@ import pandas as pd
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.model._client.model import model_version_impl
-from snowflake.ml.model._client.ops import model_ops
+from snowflake.ml.model._client.ops import model_ops, service_ops
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
@@ -19,6 +19,7 @@ class Model:
     """Model Object containing multiple versions. Mapping to SQL's MODEL object."""
 
     _model_ops: model_ops.ModelOperator
+    _service_ops: service_ops.ServiceOperator
     _model_name: sql_identifier.SqlIdentifier
 
     def __init__(self) -> None:
@@ -29,17 +30,23 @@ class Model:
         cls,
         model_ops: model_ops.ModelOperator,
         *,
+        service_ops: service_ops.ServiceOperator,
         model_name: sql_identifier.SqlIdentifier,
     ) -> "Model":
         self: "Model" = object.__new__(cls)
         self._model_ops = model_ops
+        self._service_ops = service_ops
         self._model_name = model_name
         return self
 
     def __eq__(self, __value: object) -> bool:
         if not isinstance(__value, Model):
             return False
-        return self._model_ops == __value._model_ops and self._model_name == __value._model_name
+        return (
+            self._model_ops == __value._model_ops
+            and self._service_ops == __value._service_ops
+            and self._model_name == __value._model_name
+        )
 
     @property
     def name(self) -> str:
@@ -208,6 +215,7 @@ class Model:
 
         return model_version_impl.ModelVersion._ref(
             self._model_ops,
+            service_ops=self._service_ops,
             model_name=self._model_name,
             version_name=version_id,
         )
@@ -235,6 +243,7 @@ class Model:
         return [
             model_version_impl.ModelVersion._ref(
                 self._model_ops,
+                service_ops=self._service_ops,
                 model_name=self._model_name,
                 version_name=version_name,
             )