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

snowflake/ml/data/data_connector.py

@@ -0,0 +1,133 @@
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Type, TypeVar
+
+import numpy.typing as npt
+
+from snowflake import snowpark
+from snowflake.ml._internal import telemetry
+from snowflake.ml.data import data_ingestor, data_source
+from snowflake.ml.data._internal.arrow_ingestor import ArrowIngestor as DefaultIngestor
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import tensorflow as tf
+    from torch.utils import data as torch_data
+
+    # This module can't actually depend on dataset to avoid a circular dependency
+    # Dataset -> DatasetReader -> DataConnector -!-> Dataset
+    from snowflake.ml import dataset
+
+_PROJECT = "DataConnector"
+
+DataConnectorType = TypeVar("DataConnectorType", bound="DataConnector")
+
+
+class DataConnector:
+    """Snowflake data reader which provides application integration connectors"""
+
+    def __init__(
+        self,
+        ingestor: data_ingestor.DataIngestor,
+    ) -> None:
+        self._ingestor = ingestor
+
+    @classmethod
+    def from_dataframe(cls: Type[DataConnectorType], df: snowpark.DataFrame, **kwargs: Any) -> DataConnectorType:
+        if len(df.queries["queries"]) != 1 or len(df.queries["post_actions"]) != 0:
+            raise ValueError("DataFrames with multiple queries and/or post-actions not supported")
+        source = data_source.DataFrameInfo(df.queries["queries"][0])
+        assert df._session is not None
+        ingestor = DefaultIngestor(df._session, [source])
+        return cls(ingestor, **kwargs)
+
+    @classmethod
+    def from_dataset(cls: Type[DataConnectorType], ds: "dataset.Dataset", **kwargs: Any) -> DataConnectorType:
+        dsv = ds.selected_version
+        assert dsv is not None
+        source = data_source.DatasetInfo(
+            ds.fully_qualified_name, dsv.name, dsv.url(), exclude_cols=(dsv.label_cols + dsv.exclude_cols)
+        )
+        ingestor = DefaultIngestor(ds._session, [source])
+        return cls(ingestor, **kwargs)
+
+    @property
+    def data_sources(self) -> List[data_source.DataSource]:
+        return self._ingestor.data_sources
+
+    @telemetry.send_api_usage_telemetry(
+        project=_PROJECT,
+        subproject_extractor=lambda self: type(self).__name__,
+        func_params_to_log=["batch_size", "shuffle", "drop_last_batch"],
+    )
+    def to_tf_dataset(
+        self, *, batch_size: int, shuffle: bool = False, drop_last_batch: bool = True
+    ) -> "tf.data.Dataset":
+        """Transform the Snowflake data into a ready-to-use TensorFlow tf.data.Dataset.
+
+        Args:
+            batch_size: It specifies the size of each data batch which will be
+                yield in the result datapipe
+            shuffle: It specifies whether the data will be shuffled. If True, files will be shuffled, and
+                rows in each file will also be shuffled.
+            drop_last_batch: Whether the last batch of data should be dropped. If set to be true,
+                then the last batch will get dropped if its size is smaller than the given batch_size.
+
+        Returns:
+            A tf.data.Dataset that yields batched tf.Tensors.
+        """
+        import tensorflow as tf
+
+        def generator() -> Generator[Dict[str, npt.NDArray[Any]], None, None]:
+            yield from self._ingestor.to_batches(batch_size, shuffle, drop_last_batch)
+
+        # Derive TensorFlow signature
+        first_batch = next(self._ingestor.to_batches(1, shuffle=False, drop_last_batch=False))
+        tf_signature = {
+            k: tf.TensorSpec(shape=(None,), dtype=tf.dtypes.as_dtype(v.dtype), name=k) for k, v in first_batch.items()
+        }
+
+        return tf.data.Dataset.from_generator(generator, output_signature=tf_signature)
+
+    @telemetry.send_api_usage_telemetry(
+        project=_PROJECT,
+        subproject_extractor=lambda self: type(self).__name__,
+        func_params_to_log=["batch_size", "shuffle", "drop_last_batch"],
+    )
+    def to_torch_datapipe(
+        self, *, batch_size: int, shuffle: bool = False, drop_last_batch: bool = True
+    ) -> "torch_data.IterDataPipe":  # type: ignore[type-arg]
+        """Transform the Snowflake data into a ready-to-use Pytorch datapipe.
+
+        Return a Pytorch datapipe which iterates on rows of data.
+
+        Args:
+            batch_size: It specifies the size of each data batch which will be
+                yield in the result datapipe
+            shuffle: It specifies whether the data will be shuffled. If True, files will be shuffled, and
+                rows in each file will also be shuffled.
+            drop_last_batch: Whether the last batch of data should be dropped. If set to be true,
+                then the last batch will get dropped if its size is smaller than the given batch_size.
+
+        Returns:
+            A Pytorch iterable datapipe that yield data.
+        """
+        from torch.utils.data.datapipes import iter as torch_iter
+
+        return torch_iter.IterableWrapper(  # type: ignore[no-untyped-call]
+            self._ingestor.to_batches(batch_size, shuffle, drop_last_batch)
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_PROJECT,
+        subproject_extractor=lambda self: type(self).__name__,
+        func_params_to_log=["limit"],
+    )
+    def to_pandas(self, limit: Optional[int] = None) -> "pd.DataFrame":
+        """Retrieve the Snowflake data as a Pandas DataFrame.
+
+        Args:
+            limit: If specified, the maximum number of rows to load into the DataFrame.
+
+        Returns:
+            A Pandas DataFrame.
+        """
+        return self._ingestor.to_pandas(limit)

snowflake/ml/data/data_ingestor.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Protocol, TypeVar
+
+from numpy import typing as npt
+
+from snowflake.ml.data import data_source
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+DataIngestorType = TypeVar("DataIngestorType", bound="DataIngestor")
+
+
+class DataIngestor(Protocol):
+    @property
+    def data_sources(self) -> List[data_source.DataSource]:
+        raise NotImplementedError
+
+    def to_batches(
+        self,
+        batch_size: int,
+        shuffle: bool = True,
+        drop_last_batch: bool = True,
+    ) -> Iterator[Dict[str, npt.NDArray[Any]]]:
+        raise NotImplementedError
+
+    def to_pandas(self, limit: Optional[int] = None) -> "pd.DataFrame":
+        raise NotImplementedError

snowflake/ml/data/data_source.py

@@ -0,0 +1,23 @@
+import dataclasses
+from typing import List, Optional, Union
+
+
+@dataclasses.dataclass(frozen=True)
+class DataFrameInfo:
+    """Serializable information from Snowpark DataFrames"""
+
+    sql: str
+    query_id: Optional[str] = None
+
+
+@dataclasses.dataclass(frozen=True)
+class DatasetInfo:
+    """Serializable information from SnowML Datasets"""
+
+    fully_qualified_name: str
+    version: str
+    url: Optional[str] = None
+    exclude_cols: Optional[List[str]] = None
+
+
+DataSource = Union[DataFrameInfo, DatasetInfo, str]

snowflake/ml/dataset/dataset.py

@@ -11,7 +11,6 @@ from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
 )
-from snowflake.ml._internal.lineage import data_source
 from snowflake.ml._internal.utils import (
     formatting,
     identifier,
@@ -19,6 +18,7 @@ from snowflake.ml._internal.utils import (
     snowpark_dataframe_utils,
 )
 from snowflake.ml.dataset import dataset_metadata, dataset_reader
+from snowflake.ml.lineage import lineage_node
 from snowflake.snowpark import exceptions as snowpark_exceptions, functions
 
 _PROJECT = "Dataset"
@@ -125,7 +125,7 @@ class DatasetVersion:
         return f"{self.__class__.__name__}(dataset='{self._parent.fully_qualified_name}', version='{self.name}')"
 
 
-class Dataset:
+class Dataset(lineage_node.LineageNode):
     """Represents a Snowflake Dataset which is organized into versions."""
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT)
@@ -138,18 +138,31 @@ class Dataset:
         selected_version: Optional[str] = None,
     ) -> None:
         """Initialize a lazily evaluated Dataset object"""
-        self._session = session
         self._db = database
         self._schema = schema
         self._name = name
-        self._fully_qualified_name = identifier.get_schema_level_object_identifier(database, schema, name)
+
+        super().__init__(
+            session,
+            identifier.get_schema_level_object_identifier(database, schema, name),
+            domain="dataset",
+            version=selected_version,
+        )
 
         self._version = DatasetVersion(self, selected_version) if selected_version else None
         self._reader: Optional[dataset_reader.DatasetReader] = None
 
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  name='{self._lineage_node_name}',\n"
+            f"  version='{self._version._version if self._version else None}',\n"
+            f")"
+        )
+
     @property
     def fully_qualified_name(self) -> str:
-        return self._fully_qualified_name
+        return self._lineage_node_name
 
     @property
     def selected_version(self) -> Optional[DatasetVersion]:
@@ -163,18 +176,7 @@ class Dataset:
                 original_exception=RuntimeError("No Dataset version selected."),
             )
         if self._reader is None:
-            v = self.selected_version
-            self._reader = dataset_reader.DatasetReader(
-                self._session,
-                [
-                    data_source.DataSource(
-                        fully_qualified_name=self._fully_qualified_name,
-                        version=v.name,
-                        url=v.url(),
-                        exclude_cols=(v.label_cols + v.exclude_cols),
-                    )
-                ],
-            )
+            self._reader = dataset_reader.DatasetReader.from_dataset(self, snowpark_session=self._session)
         return self._reader
 
     @staticmethod
@@ -230,9 +232,8 @@ class Dataset:
         try:
             session.sql(query).collect(statement_params=_TELEMETRY_STATEMENT_PARAMS)
             return Dataset(session, db, schema, ds_name)
-        except snowpark_exceptions.SnowparkClientException as e:
-            # Snowpark wraps the Python Connector error code in the head of the error message.
-            if e.message.startswith(dataset_errors.ERRNO_OBJECT_ALREADY_EXISTS):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code == dataset_errors.ERRNO_OBJECT_ALREADY_EXISTS:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.OBJECT_ALREADY_EXISTS,
                     original_exception=dataset_errors.DatasetExistError(
@@ -296,7 +297,7 @@ class Dataset:
         Raises:
             SnowflakeMLException: The Dataset no longer exists.
             SnowflakeMLException: The specified Dataset version already exists.
-            snowpark_exceptions.SnowparkClientException: An error occurred during Dataset creation.
+            snowpark_exceptions.SnowparkSQLException: An error occurred during Dataset creation.
 
         Note: During the generation of stage files, data casting will occur. The casting rules are as follows::
             - Data casting:
@@ -321,7 +322,8 @@ class Dataset:
                 - DateType(DATE): Not supported. A warning will be logged.
                 - VariantType(VARIANT): Not supported. A warning will be logged.
         """
-        casted_df = snowpark_dataframe_utils.cast_snowpark_dataframe(input_dataframe)
+        cast_ignore_cols = (exclude_cols or []) + (label_cols or [])
+        casted_df = snowpark_dataframe_utils.cast_snowpark_dataframe(input_dataframe, ignore_columns=cast_ignore_cols)
 
         if shuffle:
             casted_df = casted_df.order_by(functions.random())
@@ -367,19 +369,19 @@ class Dataset:
 
             return Dataset(self._session, self._db, self._schema, self._name, version)
 
-        except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(dataset_errors.ERRNO_DATASET_NOT_EXIST):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code == dataset_errors.ERRNO_DATASET_NOT_EXIST:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.NOT_FOUND,
                     original_exception=dataset_errors.DatasetNotExistError(
                         dataset_error_messages.DATASET_NOT_EXIST.format(self.fully_qualified_name)
                     ),
                 ) from e
-            elif (
-                e.message.startswith(dataset_errors.ERRNO_DATASET_VERSION_ALREADY_EXISTS)
-                or e.message.startswith(dataset_errors.ERRNO_VERSION_ALREADY_EXISTS)
-                or e.message.startswith(dataset_errors.ERRNO_FILES_ALREADY_EXISTING)
-            ):
+            elif e.sql_error_code in {
+                dataset_errors.ERRNO_DATASET_VERSION_ALREADY_EXISTS,
+                dataset_errors.ERRNO_VERSION_ALREADY_EXISTS,
+                dataset_errors.ERRNO_FILES_ALREADY_EXISTING,
+            }:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.OBJECT_ALREADY_EXISTS,
                     original_exception=dataset_errors.DatasetExistError(
@@ -435,9 +437,8 @@ class Dataset:
                 .has_column(_DATASET_VERSION_NAME_COL, allow_empty=True)
                 .validate()
             )
-        except snowpark_exceptions.SnowparkClientException as e:
-            # Snowpark wraps the Python Connector error code in the head of the error message.
-            if e.message.startswith(dataset_errors.ERRNO_OBJECT_NOT_EXIST):
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if e.sql_error_code == dataset_errors.ERRNO_OBJECT_NOT_EXIST:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.NOT_FOUND,
                     original_exception=dataset_errors.DatasetNotExistError(
@@ -459,6 +460,12 @@ class Dataset:
                 ),
             )
 
+    @staticmethod
+    def _load_from_lineage_node(session: snowpark.Session, name: str, version: str) -> "Dataset":
+        return Dataset.load(session, name).select_version(version)
+
+
+lineage_node.DOMAIN_LINEAGE_REGISTRY["dataset"] = Dataset
 
 # Utility methods
 

snowflake/ml/dataset/dataset_reader.py

@@ -1,48 +1,31 @@
-from typing import Any, List
-
-import pandas as pd
-from pyarrow import parquet as pq
+from typing import List, Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
-from snowflake.ml._internal.lineage import data_source, lineage_utils
-from snowflake.ml._internal.utils import import_utils
+from snowflake.ml._internal.lineage import lineage_utils
+from snowflake.ml.data import data_connector, data_ingestor, data_source
+from snowflake.ml.data._internal import ingestor_utils
 from snowflake.ml.fileset import snowfs
 
 _PROJECT = "Dataset"
 _SUBPROJECT = "DatasetReader"
-TARGET_FILE_SIZE = 32 * 2**20  # The max file size for data loading.
 
 
-class DatasetReader:
+class DatasetReader(data_connector.DataConnector):
     """Snowflake Dataset abstraction which provides application integration connectors"""
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
     def __init__(
         self,
-        session: snowpark.Session,
-        sources: List[data_source.DataSource],
+        ingestor: data_ingestor.DataIngestor,
+        *,
+        snowpark_session: snowpark.Session,
     ) -> None:
-        """Initialize a DatasetVersion object.
+        super().__init__(ingestor)
 
-        Args:
-            session: Snowpark Session to interact with Snowflake backend.
-            sources: Data sources to read from.
-
-        Raises:
-            ValueError: `sources` arg was empty or null
-        """
-        if not sources:
-            raise ValueError("Invalid input: empty `sources` list not allowed")
-        self._session = session
-        self._sources = sources
-        self._fs: snowfs.SnowFileSystem = snowfs.SnowFileSystem(
-            snowpark_session=self._session,
-            cache_type="bytes",
-            block_size=2 * TARGET_FILE_SIZE,
-        )
-
-        self._files: List[str] = []
+        self._session: snowpark.Session = snowpark_session
+        self._fs: snowfs.SnowFileSystem = ingestor_utils.get_dataset_filesystem(self._session)
+        self._files: Optional[List[str]] = None
 
     def _list_files(self) -> List[str]:
         """Private helper function that lists all files in this DatasetVersion and caches the results."""
@@ -50,18 +33,14 @@ class DatasetReader:
             return self._files
 
         files: List[str] = []
-        for source in self._sources:
-            # Sort within each source for consistent ordering
-            files.extend(sorted(self._fs.ls(source.url)))  # type: ignore[arg-type]
+        for source in self.data_sources:
+            assert isinstance(source, data_source.DatasetInfo)
+            files.extend(ingestor_utils.get_dataset_files(self._session, source, filesystem=self._fs))
         files.sort()
 
         self._files = files
         return self._files
 
-    @property
-    def data_sources(self) -> List[data_source.DataSource]:
-        return self._sources
-
     @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
     def files(self) -> List[str]:
         """Get the list of remote file paths for the current DatasetVersion.
@@ -85,76 +64,6 @@ class DatasetReader:
         """Return an fsspec FileSystem which can be used to load the DatasetVersion's `files()`"""
         return self._fs
 
-    @telemetry.send_api_usage_telemetry(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        func_params_to_log=["batch_size", "shuffle", "drop_last_batch"],
-    )
-    def to_torch_datapipe(self, *, batch_size: int, shuffle: bool = False, drop_last_batch: bool = True) -> Any:
-        """Transform the Snowflake data into a ready-to-use Pytorch datapipe.
-
-        Return a Pytorch datapipe which iterates on rows of data.
-
-        Args:
-            batch_size: It specifies the size of each data batch which will be
-                yield in the result datapipe
-            shuffle: It specifies whether the data will be shuffled. If True, files will be shuffled, and
-                rows in each file will also be shuffled.
-            drop_last_batch: Whether the last batch of data should be dropped. If set to be true,
-                then the last batch will get dropped if its size is smaller than the given batch_size.
-
-        Returns:
-            A Pytorch iterable datapipe that yield data.
-
-        Examples:
-        >>> dp = dataset.to_torch_datapipe(batch_size=1)
-        >>> for data in dp:
-        >>>     print(data)
-        ----
-        {'_COL_1':[10]}
-        """
-        IterableWrapper, _ = import_utils.import_or_get_dummy("torchdata.datapipes.iter.IterableWrapper")
-        torch_datapipe_module, _ = import_utils.import_or_get_dummy("snowflake.ml.fileset.torch_datapipe")
-
-        self._fs.optimize_read(self._list_files())
-
-        input_dp = IterableWrapper(self._list_files())
-        return torch_datapipe_module.ReadAndParseParquet(input_dp, self._fs, batch_size, shuffle, drop_last_batch)
-
-    @telemetry.send_api_usage_telemetry(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        func_params_to_log=["batch_size", "shuffle", "drop_last_batch"],
-    )
-    def to_tf_dataset(self, *, batch_size: int, shuffle: bool = False, drop_last_batch: bool = True) -> Any:
-        """Transform the Snowflake data into a ready-to-use TensorFlow tf.data.Dataset.
-
-        Args:
-            batch_size: It specifies the size of each data batch which will be
-                yield in the result datapipe
-            shuffle: It specifies whether the data will be shuffled. If True, files will be shuffled, and
-                rows in each file will also be shuffled.
-            drop_last_batch: Whether the last batch of data should be dropped. If set to be true,
-                then the last batch will get dropped if its size is smaller than the given batch_size.
-
-        Returns:
-            A tf.data.Dataset that yields batched tf.Tensors.
-
-        Examples:
-        >>> dp = dataset.to_tf_dataset(batch_size=1)
-        >>> for data in dp:
-        >>>     print(data)
-        ----
-        {'_COL_1': <tf.Tensor: shape=(1,), dtype=int64, numpy=[10]>}
-        """
-        tf_dataset_module, _ = import_utils.import_or_get_dummy("snowflake.ml.fileset.tf_dataset")
-
-        self._fs.optimize_read(self._list_files())
-
-        return tf_dataset_module.read_and_parse_parquet(
-            self._list_files(), self._fs, batch_size, shuffle, drop_last_batch
-        )
-
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
         subproject=_SUBPROJECT,
@@ -177,7 +86,8 @@ class DatasetReader:
         """
         file_path_pattern = ".*data_.*[.]parquet"
         dfs: List[snowpark.DataFrame] = []
-        for source in self._sources:
+        for source in self.data_sources:
+            assert isinstance(source, data_source.DatasetInfo) and source.url is not None
             df = self._session.read.option("pattern", file_path_pattern).parquet(source.url)
             if only_feature_cols and source.exclude_cols:
                 df = df.drop(source.exclude_cols)
@@ -186,14 +96,4 @@ class DatasetReader:
         combined_df = dfs[0]
         for df in dfs[1:]:
             combined_df = combined_df.union_all_by_name(df)
-        return lineage_utils.patch_dataframe(combined_df, data_sources=self._sources, inplace=True)
-
-    @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
-    def to_pandas(self) -> pd.DataFrame:
-        """Retrieve the DatasetVersion contents as a Pandas Dataframe"""
-        files = self._list_files()
-        if not files:
-            return pd.DataFrame()  # Return empty DataFrame
-        self._fs.optimize_read(files)
-        pd_ds = pq.ParquetDataset(files, filesystem=self._fs)
-        return pd_ds.read_pandas().to_pandas()
+        return lineage_utils.patch_dataframe(combined_df, data_sources=self.data_sources, inplace=True)

snowflake/ml/feature_store/access_manager.py

@@ -273,7 +273,13 @@ def setup_feature_store(
     assert current_role is not None  # to make mypy happy
     try:
         session.use_role(producer_role)
-        fs = FeatureStore(session, database, schema, warehouse, creation_mode=CreationMode.CREATE_IF_NOT_EXIST)
+        fs = FeatureStore(
+            session,
+            database,
+            schema,
+            default_warehouse=warehouse,
+            creation_mode=CreationMode.CREATE_IF_NOT_EXIST,
+        )
     finally:
         session.use_role(current_role)
 

snowflake/ml/feature_store/entity.py

@@ -22,7 +22,7 @@ class Entity:
     It can also be used for FeatureView search and lineage tracking.
     """
 
-    def __init__(self, name: str, join_keys: List[str], desc: str = "") -> None:
+    def __init__(self, name: str, join_keys: List[str], *, desc: str = "") -> None:
         """
         Creates an Entity instance.
 
@@ -30,6 +30,23 @@ class Entity:
             name: name of the Entity.
             join_keys: join keys associated with a FeatureView, used for feature retrieval.
             desc: description of the Entity.
+
+        Example::
+
+            >>> fs = FeatureStore(...)
+            >>> e_1 = Entity(
+            ...     name="my_entity",
+            ...     join_keys=['col_1'],
+            ...     desc='My first entity.'
+            ... )
+            >>> fs.register_entity(e_1)
+            >>> fs.list_entities().show()
+            -----------------------------------------------------------
+            |"NAME"     |"JOIN_KEYS"  |"DESC"            |"OWNER"     |
+            -----------------------------------------------------------
+            |MY_ENTITY  |["COL_1"]    |My first entity.  |REGTEST_RL  |
+            -----------------------------------------------------------
+
         """
         self._validate(name, join_keys)
 
@@ -65,7 +82,7 @@ class Entity:
 
     @staticmethod
     def _construct_entity(name: str, join_keys: List[str], desc: str, owner: str) -> "Entity":
-        e = Entity(name, join_keys, desc)
+        e = Entity(name, join_keys, desc=desc)
         e.owner = owner
         return e
 

snowflake/ml/feature_store/examples/citibike_trip_features/entities.py

@@ -0,0 +1,20 @@
+from typing import List
+
+from snowflake.ml.feature_store import Entity
+
+end_station_id = Entity(
+    name="end_station_id",
+    join_keys=["end_station_id"],
+    desc="The id of an end station.",
+)
+
+trip_id = Entity(
+    name="trip_id",
+    join_keys=["trip_id"],
+    desc="The id of a trip.",
+)
+
+
+# This will be invoked by example_helper.py. Do not change function name.
+def get_all_entities() -> List[Entity]:
+    return [end_station_id, trip_id]

snowflake/ml/feature_store/examples/citibike_trip_features/features/station_feature.py

@@ -0,0 +1,31 @@
+from typing import List
+
+from snowflake.ml.feature_store import FeatureView
+from snowflake.ml.feature_store.examples.citibike_trip_features.entities import (
+    end_station_id,
+)
+from snowflake.snowpark import DataFrame, Session
+
+
+# This function will be invoked by example_helper.py. Do not change the name.
+def create_draft_feature_view(session: Session, source_dfs: List[DataFrame], source_tables: List[str]) -> FeatureView:
+    """Create a feature view about trip station."""
+    query = session.sql(
+        f"""
+        select
+            end_station_id,
+            count(end_station_id) as f_count_1d,
+            avg(end_station_latitude) as f_avg_latitude_1d,
+            avg(end_station_longitude) as f_avg_longtitude_1d
+        from {source_tables[0]}
+        group by end_station_id
+        """
+    )
+
+    return FeatureView(
+        name="f_station_1d",  # name of feature view
+        entities=[end_station_id],  # entities
+        feature_df=query,  # definition query
+        refresh_freq="1d",  # refresh frequency. '1d' means it refreshes everyday
+        desc="Station features refreshed every day.",
+    )

snowflake/ml/feature_store/examples/citibike_trip_features/features/trip_feature.py

@@ -0,0 +1,24 @@
+from typing import List
+
+from snowflake.ml.feature_store import FeatureView
+from snowflake.ml.feature_store.examples.citibike_trip_features.entities import trip_id
+from snowflake.snowpark import DataFrame, Session, functions as F
+
+
+# This function will be invoked by example_helper.py. Do not change the name.
+def create_draft_feature_view(session: Session, source_dfs: List[DataFrame], source_tables: List[str]) -> FeatureView:
+    """Create a feature view about trip."""
+    feature_df = source_dfs[0].select(
+        "trip_id",
+        F.col("birth_year").alias("f_birth_year"),
+        F.col("gender").alias("f_gender"),
+        F.col("bikeid").alias("f_bikeid"),
+    )
+
+    return FeatureView(
+        name="f_trip",  # name of feature view
+        entities=[trip_id],  # entities
+        feature_df=feature_df,  # definition query
+        refresh_freq=None,  # refresh frequency. None indicates it never refresh
+        desc="Static trip features",
+    )

snowflake/ml/feature_store/examples/citibike_trip_features/source.yaml

@@ -0,0 +1,4 @@
+---
+source_data: citibike_trips
+label_columns: tripduration
+add_id_column: trip_id