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

snowflake/ml/feature_store/examples/airline_features/features/plane_features.py

@@ -0,0 +1,31 @@
+from typing import List
+
+from snowflake.ml.feature_store import FeatureView
+from snowflake.ml.feature_store.examples.airline_features.entities import plane_entity
+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 airplane model."""
+    query = session.sql(
+        """
+        select
+           PLANE_MODEL,
+           SEATING_CAPACITY
+        from
+            PLANE_MODEL_ATTRIBUTES
+        """
+    )
+
+    return FeatureView(
+        name="f_plane",  # name of feature view
+        entities=[plane_entity],  # entities
+        feature_df=query,  # definition query
+        refresh_freq=None,  # refresh frequency
+        desc="Plane features never refresh.",
+    ).attach_feature_desc(
+        {
+            "SEATING_CAPACITY": "The seating capacity of a plane.",
+        }
+    )

snowflake/ml/feature_store/examples/airline_features/features/weather_features.py

@@ -0,0 +1,42 @@
+from typing import List
+
+from snowflake.ml.feature_store import FeatureView
+from snowflake.ml.feature_store.examples.airline_features.entities import zipcode_entity
+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 airport weather."""
+    query = session.sql(
+        """
+        select
+            DATETIME_UTC AS TS,
+            AIRPORT_ZIP_CODE,
+            sum(RAIN_MM_H) over (
+                partition by AIRPORT_ZIP_CODE
+                order by DATETIME_UTC
+                range between interval '30 minutes' preceding and current row
+            ) RAIN_SUM_30M,
+            sum(RAIN_MM_H) over (
+                partition by AIRPORT_ZIP_CODE
+                order by DATETIME_UTC
+                range between interval '1 day' preceding and current row
+            ) RAIN_SUM_60M
+        from AIRPORT_WEATHER_STATION
+        """
+    )
+
+    return FeatureView(
+        name="f_weather",  # name of feature view
+        entities=[zipcode_entity],  # entities
+        feature_df=query,  # definition query
+        timestamp_col="TS",  # timestamp column
+        refresh_freq="1d",  # refresh frequency
+        desc="Airport weather features refreshed every day.",
+    ).attach_feature_desc(
+        {
+            "RAIN_SUM_30M": "The sum of rain fall over past 30 minutes for one zipcode.",
+            "RAIN_SUM_60M": "The sum of rain fall over past 1 day for one zipcode.",
+        }
+    )

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

@@ -0,0 +1,7 @@
+---
+source_data: airline
+label_columns: DEPARTING_DELAY
+timestamp_column: SCHEDULED_DEPARTURE_UTC
+desc: Features using synthetic airline data to predict the departing delay.
+model_category: classification
+training_spine_table: US_FLIGHT_SCHEDULES

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

@@ -14,18 +14,24 @@ def create_draft_feature_view(session: Session, source_dfs: List[DataFrame], sou
         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
+            count(end_station_id) as f_count,
+            avg(end_station_latitude) as f_avg_latitude,
+            avg(end_station_longitude) as f_avg_longtitude
         from {source_tables[0]}
         group by end_station_id
         """
     )
 
     return FeatureView(
-        name="f_station_1d",  # name of feature view
+        name="f_station",  # 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.",
+    ).attach_feature_desc(
+        {
+            "f_count": "How many times this station appears in 1 day.",
+            "f_avg_latitude": "Averaged latitude of a station.",
+            "f_avg_longtitude": "Averaged longtitude of a station.",
+        }
     )

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

@@ -21,4 +21,10 @@ def create_draft_feature_view(session: Session, source_dfs: List[DataFrame], sou
         feature_df=feature_df,  # definition query
         refresh_freq=None,  # refresh frequency. None indicates it never refresh
         desc="Static trip features",
+    ).attach_feature_desc(
+        {
+            "f_birth_year": "The birth year of a trip passenger.",
+            "f_gender": "The gender of a trip passenger.",
+            "f_bikeid": "The bike id of a trip passenger.",
+        }
     )

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

@@ -1,4 +1,7 @@
 ---
 source_data: citibike_trips
+training_spine_table: citibike_trips
 label_columns: tripduration
 add_id_column: trip_id
+desc: Features using citibike trip data trying to predict the duration of a trip.
+model_category: regression

snowflake/ml/feature_store/examples/example_helper.py

@@ -10,7 +10,7 @@ import yaml
 from snowflake.ml._internal.utils import identifier, sql_identifier
 from snowflake.ml.feature_store import Entity, FeatureView  # type: ignore[attr-defined]
 from snowflake.snowpark import DataFrame, Session, functions as F
-from snowflake.snowpark.types import TimestampType
+from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -28,6 +28,9 @@ class ExampleHelper:
         self._session = session
         self._database_name = database_name
         self._dataset_schema = dataset_schema
+        self._clear()
+
+    def _clear(self) -> None:
         self._selected_example = None
         self._source_tables: List[str] = []
         self._source_dfs: List[DataFrame] = []
@@ -36,15 +39,18 @@ class ExampleHelper:
         self._timestamp_column: Optional[sql_identifier.SqlIdentifier] = None
         self._epoch_to_timestamp_cols: List[str] = []
         self._add_id_column: Optional[sql_identifier.SqlIdentifier] = None
+        self._training_spine_table: str = ""
 
-    def list_examples(self) -> List[str]:
-        """Return a list of examples."""
+    def list_examples(self) -> Optional[DataFrame]:
+        """Return a dataframe object about descriptions of all examples."""
         root_dir = Path(__file__).parent
-        result = []
+        rows = []
         for f_name in os.listdir(root_dir):
             if os.path.isdir(os.path.join(root_dir, f_name)) and f_name[0].isalpha() and f_name != "source_data":
-                result.append(f_name)
-        return result
+                source_file_path = root_dir.joinpath(f"{f_name}/source.yaml")
+                source_dict = self._read_yaml(str(source_file_path))
+                rows.append((f_name, source_dict["model_category"], source_dict["desc"], source_dict["label_columns"]))
+        return self._session.create_dataframe(rows, schema=["NAME", "MODEL_CATEGORY", "DESC", "LABEL_COLS"])
 
     def load_draft_feature_views(self) -> List[FeatureView]:
         """Return all feature views in an example.
@@ -101,7 +107,7 @@ class ExampleHelper:
             """
         ).collect()
 
-    def _load_csv(self, schema_dict: Dict[str, str], destination_table: str, temp_stage_name: str) -> None:
+    def _load_csv(self, schema_dict: Dict[str, str], temp_stage_name: str) -> List[str]:
         # create temp file format
         file_format_name = f"{self._database_name}.{self._dataset_schema}.feature_store_temp_format"
         format_str = ""
@@ -116,6 +122,8 @@ class ExampleHelper:
             cols_type_str = (
                 f"{self._add_id_column.resolved()} number autoincrement start 1 increment 1, " + cols_type_str
             )
+
+        destination_table = f"{self._database_name}.{self._dataset_schema}.{schema_dict['destination_table_name']}"
         self._session.sql(
             f"""
             create or replace table {destination_table} ({cols_type_str})
@@ -132,25 +140,50 @@ class ExampleHelper:
             """
         ).collect()
 
-    def _load_parquet(self, schema_dict: Dict[str, str], destination_table: str, temp_stage_name: str) -> None:
+        return [destination_table]
+
+    def _load_parquet(self, schema_dict: Dict[str, str], temp_stage_name: str) -> List[str]:
         regex_pattern = schema_dict["load_files_pattern"]
         all_files = self._session.sql(f"list @{temp_stage_name}").collect()
         filtered_files = [item["name"] for item in all_files if re.match(regex_pattern, item["name"])]
-        assert len(filtered_files) == 1, "Current code only works for one file"
-        file_name = filtered_files[0].rsplit("/", 1)[-1]
+        file_count = len(filtered_files)
+        result = []
+
+        for file in filtered_files:
+            file_name = file.rsplit("/", 1)[-1]
 
-        df = self._session.read.parquet(f"@{temp_stage_name}/{file_name}")
-        for old_col_name in df.columns:
-            df = df.with_column_renamed(old_col_name, identifier.get_unescaped_names(old_col_name))
+            df = self._session.read.parquet(f"@{temp_stage_name}/{file_name}")
+            for old_col_name in df.columns:
+                df = df.with_column_renamed(old_col_name, identifier.get_unescaped_names(old_col_name))
 
-        for ts_col in self._epoch_to_timestamp_cols:
-            if "timestamp" != dict(df.dtypes)[ts_col]:
-                df = df.with_column(f"{ts_col}_NEW", F.cast(df[ts_col] / 1000000, TimestampType()))
-                df = df.drop(ts_col).rename(f"{ts_col}_NEW", ts_col)
+            # convert timestamp column to ntz
+            for name, type in dict(df.dtypes).items():
+                if type == "timestamp":
+                    df = df.with_column(name, F.to_timestamp_ntz(name))
 
-        df.write.mode("overwrite").save_as_table(destination_table)
+            # convert epoch column to ntz timestamp
+            for ts_col in self._epoch_to_timestamp_cols:
+                if "timestamp" != dict(df.dtypes)[ts_col]:
+                    df = df.with_column(ts_col, F.cast(df[ts_col] / 1000000, TimestampType(TimestampTimeZone.NTZ)))
 
-    def _load_source_data(self, schema_yaml_file: str) -> str:
+            if self._add_id_column:
+                df = df.withColumn(self._add_id_column, F.monotonically_increasing_id())
+
+            if file_count == 1:
+                dest_table_name = (
+                    f"{self._database_name}.{self._dataset_schema}.{schema_dict['destination_table_name']}"
+                )
+            else:
+                regex_pattern = schema_dict["destination_table_name"]
+                dest_table_name = re.match(regex_pattern, file_name).group("table_name")  # type: ignore[union-attr]
+                dest_table_name = f"{self._database_name}.{self._dataset_schema}.{dest_table_name}"
+
+            df.write.mode("overwrite").save_as_table(dest_table_name)
+            result.append(dest_table_name)
+
+        return result
+
+    def _load_source_data(self, schema_yaml_file: str) -> List[str]:
         """Parse a yaml schema file and load data into Snowflake.
 
         Args:
@@ -162,7 +195,6 @@ class ExampleHelper:
         # load schema file
         schema_dict = self._read_yaml(schema_yaml_file)
         temp_stage_name = f"{self._database_name}.{self._dataset_schema}.feature_store_temp_stage"
-        destination_table = f"{self._database_name}.{self._dataset_schema}.{schema_dict['destination_table_name']}"
 
         # create a temp stage from S3 URL
         self._session.sql(f"create or replace stage {temp_stage_name} url = '{schema_dict['s3_url']}'").collect()
@@ -170,11 +202,9 @@ class ExampleHelper:
         # load csv or parquet
         # TODO: this could be more flexible and robust.
         if "parquet" in schema_dict["load_files_pattern"]:
-            self._load_parquet(schema_dict, destination_table, temp_stage_name)
+            return self._load_parquet(schema_dict, temp_stage_name)
         else:
-            self._load_csv(schema_dict, destination_table, temp_stage_name)
-
-        return destination_table
+            return self._load_csv(schema_dict, temp_stage_name)
 
     def load_example(self, example_name: str) -> List[str]:
         """Select the active example and load its datasets to Snowflake.
@@ -186,6 +216,7 @@ class ExampleHelper:
         Returns:
             Returns a list of table names with populated datasets.
         """
+        self._clear()
         self._selected_example = example_name  # type: ignore[assignment]
 
         # load source yaml file
@@ -195,7 +226,7 @@ class ExampleHelper:
         self._source_tables = []
         self._source_dfs = []
 
-        source_ymal_data = source_dict["source_data"]
+        source_yaml_data = source_dict["source_data"]
         if "excluded_columns" in source_dict:
             self._excluded_columns = sql_identifier.to_sql_identifiers(source_dict["excluded_columns"].split(","))
         if "label_columns" in source_dict:
@@ -206,8 +237,11 @@ class ExampleHelper:
             self._epoch_to_timestamp_cols = source_dict["epoch_to_timestamp_cols"].split(",")
         if "add_id_column" in source_dict:
             self._add_id_column = sql_identifier.SqlIdentifier(source_dict["add_id_column"])
+        self._training_spine_table = (
+            f"{self._database_name}.{self._dataset_schema}.{source_dict['training_spine_table']}"
+        )
 
-        return self.load_source_data(source_ymal_data)
+        return self.load_source_data(source_yaml_data)
 
     def load_source_data(self, source_data_name: str) -> List[str]:
         """Load source data into Snowflake.
@@ -220,11 +254,12 @@ class ExampleHelper:
         """
         root_dir = Path(__file__).parent
         schema_file = root_dir.joinpath(f"source_data/{source_data_name}.yaml")
-        destination_table = self._load_source_data(str(schema_file))
-        source_df = self._session.table(destination_table)
-        self._source_tables.append(destination_table)
-        self._source_dfs.append(source_df)
-        logger.info(f"source data {source_data_name} has been successfully loaded into table {destination_table}.")
+        destination_tables = self._load_source_data(str(schema_file))
+        for dest_table in destination_tables:
+            source_df = self._session.table(dest_table)
+            self._source_tables.append(dest_table)
+            self._source_dfs.append(source_df)
+            logger.info(f"{dest_table} has been created successfully.")
         return self._source_tables
 
     def get_current_schema(self) -> str:
@@ -238,3 +273,6 @@ class ExampleHelper:
 
     def get_training_data_timestamp_col(self) -> Optional[str]:
         return self._timestamp_column.resolved() if self._timestamp_column is not None else None
+
+    def get_training_spine_table(self) -> str:
+        return self._training_spine_table

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

@@ -2,11 +2,11 @@ from typing import List
 
 from snowflake.ml.feature_store import Entity
 
-trip_pickup = Entity(name="TRIP_PICKUP", join_keys=["PULOCATIONID"], desc="Trip pickup entity.")
+trip_id = Entity(name="TRIP_ID", join_keys=["TRIP_ID"], desc="Trip id.")
 
-trip_dropoff = Entity(name="TRIP_DROPOFF", join_keys=["DOLOCATIONID"], desc="Trip dropoff entity.")
+location_id = Entity(name="DOLOCATIONID", join_keys=["DOLOCATIONID"], desc="Drop off location id.")
 
 
 # This will be invoked by example_helper.py. Do not change function name.
 def get_all_entities() -> List[Entity]:
-    return [trip_pickup, trip_dropoff]
+    return [trip_id, location_id]

snowflake/ml/feature_store/examples/new_york_taxi_features/features/{dropoff_features.py → location_features.py}

@@ -2,7 +2,7 @@ from typing import List
 
 from snowflake.ml.feature_store import FeatureView
 from snowflake.ml.feature_store.examples.new_york_taxi_features.entities import (
-    trip_dropoff,
+    location_id,
 )
 from snowflake.snowpark import DataFrame, Session
 
@@ -15,25 +15,30 @@ def create_draft_feature_view(session: Session, source_dfs: List[DataFrame], sou
         select
             TPEP_DROPOFF_DATETIME as TS,
             DOLOCATIONID,
-            count(FARE_AMOUNT) over (
+            avg(FARE_AMOUNT) over (
                 partition by DOLOCATIONID
                 order by TPEP_DROPOFF_DATETIME
                 range between interval '1 hours' preceding and current row
-            ) TRIP_COUNT_1H,
-            count(FARE_AMOUNT) over (
+            ) AVG_FARE_1H,
+            avg(FARE_AMOUNT) over (
                 partition by DOLOCATIONID
                 order by TPEP_DROPOFF_DATETIME
-                range between interval '5 hours' preceding and current row
-            ) TRIP_COUNT_5H
+                range between interval '10 hours' preceding and current row
+            ) AVG_FARE_10h
         from {source_tables[0]}
     """
     )
 
     return FeatureView(
-        name="f_trip_dropoff",  # name of feature view
-        entities=[trip_dropoff],  # entities
+        name="f_location",  # name of feature view
+        entities=[location_id],  # entities
         feature_df=feature_df,  # definition query
         refresh_freq="12h",  # the frequency this feature view re-compute
         timestamp_col="TS",  # timestamp column. Used when generate training data
-        desc="Managed feature view trip dropoff refreshed every 12 hours.",
+        desc="Features aggregated by location id and refreshed every 12 hours.",
+    ).attach_feature_desc(
+        {
+            "AVG_FARE_1H": "Averaged fare in past 1 hour window aggregated by location.",
+            "AVG_FARE_10H": "Averaged fare in past 10 hours aggregated by location.",
+        }
     )

snowflake/ml/feature_store/examples/new_york_taxi_features/features/trip_features.py

@@ -0,0 +1,36 @@
+from typing import List
+
+from snowflake.ml.feature_store import FeatureView
+from snowflake.ml.feature_store.examples.new_york_taxi_features.entities import trip_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 draft feature view."""
+    feature_df = session.sql(
+        f"""
+        select
+            TRIP_ID,
+            PASSENGER_COUNT,
+            TRIP_DISTANCE,
+            FARE_AMOUNT
+        from
+            {source_tables[0]}
+        """
+    )
+
+    return FeatureView(
+        name="f_trip",  # name of feature view
+        entities=[trip_id],  # entities
+        feature_df=feature_df,  # definition query
+        refresh_freq="1d",  # the frequency this feature view re-compute
+        timestamp_col=None,  # timestamp column. Used when generate training data
+        desc="Features per trip refreshed every day.",
+    ).attach_feature_desc(
+        {
+            "PASSENGER_COUNT": "The count of passenger of a trip.",
+            "TRIP_DISTANCE": "The distance of a trip.",
+            "FARE_AMOUNT": "The fare of a trip.",
+        }
+    )

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

@@ -1,5 +1,9 @@
 ---
 source_data: nyc_yellow_trips
-label_columns: FARE_AMOUNT
+training_spine_table: nyc_yellow_trips
+label_columns: TOTAL_AMOUNT
+add_id_column: TRIP_ID
 timestamp_column: TPEP_PICKUP_DATETIME
 epoch_to_timestamp_cols: TPEP_PICKUP_DATETIME,TPEP_DROPOFF_DATETIME
+desc: Features using taxi trip data trying to predict the total fare of a trip.
+model_category: regression

snowflake/ml/feature_store/examples/source_data/airline.yaml

@@ -0,0 +1,4 @@
+---
+s3_url: s3://sfquickstarts/misc/demos/airline/
+load_files_pattern: .*[.]parquet
+destination_table_name: (?P<table_name>.*)_0_0_0[.]snappy[.]parquet

snowflake/ml/feature_store/examples/source_data/citibike_trips.yaml

@@ -1,7 +1,7 @@
 ---
 s3_url: s3://snowflake-workshop-lab/citibike-trips-csv/
 destination_table_name: citibike_trips
-load_files_pattern: trips_2013_6_4_0.csv.gz
+load_files_pattern: .*trips_2013_6_.*[.]csv[.]gz
 format:
   type: csv
   compression: auto

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

@@ -2,13 +2,13 @@ from typing import List
 
 from snowflake.ml.feature_store import Entity
 
-wine_entity = Entity(
+wine_id = Entity(
     name="WINE",
     join_keys=["WINE_ID"],
-    desc="Wine ID column.",
+    desc="Wine ID.",
 )
 
 
 # This will be invoked by example_helper.py. Do not change function name.
 def get_all_entities() -> List[Entity]:
-    return [wine_entity]
+    return [wine_id]

snowflake/ml/feature_store/examples/wine_quality_features/features/managed_wine_features.py

@@ -1,9 +1,7 @@
 from typing import List
 
 from snowflake.ml.feature_store import FeatureView
-from snowflake.ml.feature_store.examples.wine_quality_features.entities import (
-    wine_entity,
-)
+from snowflake.ml.feature_store.examples.wine_quality_features.entities import wine_id
 from snowflake.snowpark import DataFrame, Session, functions as F
 
 
@@ -17,13 +15,22 @@ def create_draft_feature_view(session: Session, source_dfs: List[DataFrame], sou
         "CHLORIDES",
         "TOTAL_SULFUR_DIOXIDE",
         "PH",
-        (F.col("FIXED_ACIDITY") * F.col("CITRIC_ACID")).alias("MY_NEW_FEATURE"),
+        (F.col("FIXED_ACIDITY") * F.col("CITRIC_ACID")).alias("HYBRID_ACID"),
     )
 
     return FeatureView(
         name="WINE_FEATURES",  # name of feature view
-        entities=[wine_entity],  # entities
+        entities=[wine_id],  # entities
         feature_df=feature_df,  # definition dataframe
         refresh_freq="1d",  # refresh frequency. '1d' means it refreshes everyday
-        desc="Managed feature view about wine quality which refreshes everyday.",
+        desc="Managed features about wine quality which refreshes everyday.",
+    ).attach_feature_desc(
+        {
+            "FIXED_ACIDITY": "Fixed acidity.",
+            "CITRIC_ACID": "Citric acid.",
+            "CHLORIDES": "Chlorides",
+            "TOTAL_SULFUR_DIOXIDE": "Total sulfur dioxide.",
+            "PH": "PH.",
+            "HYBRID_ACID": "Hybrid acid generated by a production of fixed and citric acid.",
+        }
     )

snowflake/ml/feature_store/examples/wine_quality_features/features/static_wine_features.py

@@ -1,9 +1,7 @@
 from typing import List
 
 from snowflake.ml.feature_store import FeatureView
-from snowflake.ml.feature_store.examples.wine_quality_features.entities import (
-    wine_entity,
-)
+from snowflake.ml.feature_store.examples.wine_quality_features.entities import wine_id
 from snowflake.snowpark import DataFrame, Session
 
 
@@ -14,8 +12,13 @@ def create_draft_feature_view(session: Session, source_dfs: List[DataFrame], sou
 
     return FeatureView(
         name="EXTRA_WINE_FEATURES",  # name of feature view
-        entities=[wine_entity],  # entities
+        entities=[wine_id],  # entities
         feature_df=feature_df,  # feature dataframe
         refresh_freq=None,  # refresh frequency. None means it never refresh
-        desc="Static feature view about wine quality which never refresh.",
+        desc="Static features about wine quality which never refresh.",
+    ).attach_feature_desc(
+        {
+            "SULPHATES": "Sulphates.",
+            "ALCOHOL": "Alcohol.",
+        }
     )

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

@@ -1,5 +1,8 @@
 ---
 source_data: winequality_red
+training_spine_table: winedata
 add_id_column: wine_id
 label_columns: quality
 excluded_columns: wine_id
+desc: Features using wine quality data trying to predict the quality of wine.
+model_category: regression