snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

snowflake/ml/feature_store/entity.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Optional
 
 from snowflake.ml._internal.utils.sql_identifier import (
     SqlIdentifier,
@@ -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.
 
@@ -51,11 +51,11 @@ class Entity:
         self._validate(name, join_keys)
 
         self.name: SqlIdentifier = SqlIdentifier(name)
-        self.join_keys: List[SqlIdentifier] = to_sql_identifiers(join_keys)
+        self.join_keys: list[SqlIdentifier] = to_sql_identifiers(join_keys)
         self.owner: Optional[str] = None
         self.desc: str = desc
 
-    def _validate(self, name: str, join_keys: List[str]) -> None:
+    def _validate(self, name: str, join_keys: list[str]) -> None:
         if len(name) > _ENTITY_NAME_LENGTH_LIMIT:
             raise ValueError(f"Entity name `{name}` exceeds maximum length: {_ENTITY_NAME_LENGTH_LIMIT}")
         if _FEATURE_VIEW_ENTITY_TAG_DELIMITER in name:
@@ -73,7 +73,7 @@ class Entity:
             if len(k) > _ENTITY_JOIN_KEY_LENGTH_LIMIT:
                 raise ValueError(f"Join key: {k} exceeds length limit {_ENTITY_JOIN_KEY_LENGTH_LIMIT}.")
 
-    def _to_dict(self) -> Dict[str, str]:
+    def _to_dict(self) -> dict[str, str]:
         entity_dict = self.__dict__.copy()
         for k, v in entity_dict.items():
             if isinstance(v, SqlIdentifier):
@@ -81,7 +81,7 @@ class Entity:
         return entity_dict
 
     @staticmethod
-    def _construct_entity(name: str, join_keys: List[str], desc: str, owner: str) -> "Entity":
+    def _construct_entity(name: str, join_keys: list[str], desc: str, owner: str) -> "Entity":
         e = Entity(name, join_keys, desc=desc)
         e.owner = owner
         return e

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

@@ -1,5 +1,3 @@
-from typing import List
-
 from snowflake.ml.feature_store import Entity
 
 zipcode_entity = Entity(
@@ -12,5 +10,5 @@ plane_entity = Entity(name="PLANE_MODEL", join_keys=["PLANE_MODEL"], desc="The m
 
 
 # This will be invoked by example_helper.py. Do not change function name.
-def get_all_entities() -> List[Entity]:
+def get_all_entities() -> list[Entity]:
     return [zipcode_entity, plane_entity]

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

@@ -1,5 +1,3 @@
-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
@@ -7,7 +5,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a feature view about airplane model."""
     query = session.sql(

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

@@ -1,5 +1,3 @@
-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
@@ -7,7 +5,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a feature view about airport weather."""
     query = session.sql(

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

@@ -1,5 +1,3 @@
-from typing import List
-
 from snowflake.ml.feature_store import Entity
 
 end_station_id = Entity(
@@ -16,5 +14,5 @@ trip_id = Entity(
 
 
 # This will be invoked by example_helper.py. Do not change function name.
-def get_all_entities() -> List[Entity]:
+def get_all_entities() -> list[Entity]:
     return [end_station_id, trip_id]

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

@@ -1,5 +1,3 @@
-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,
@@ -9,7 +7,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a feature view about trip station."""
     query = session.sql(

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

@@ -1,5 +1,3 @@
-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
@@ -7,7 +5,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a feature view about trip."""
     feature_df = source_dfs[0].select(

snowflake/ml/feature_store/examples/example_helper.py

@@ -3,7 +3,7 @@ import logging
 import os
 import re
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 import yaml
 
@@ -32,12 +32,12 @@ class ExampleHelper:
 
     def _clear(self) -> None:
         self._selected_example = None
-        self._source_tables: List[str] = []
-        self._source_dfs: List[DataFrame] = []
-        self._excluded_columns: List[sql_identifier.SqlIdentifier] = []
-        self._label_columns: List[sql_identifier.SqlIdentifier] = []
+        self._source_tables: list[str] = []
+        self._source_dfs: list[DataFrame] = []
+        self._excluded_columns: list[sql_identifier.SqlIdentifier] = []
+        self._label_columns: list[sql_identifier.SqlIdentifier] = []
         self._timestamp_column: Optional[sql_identifier.SqlIdentifier] = None
-        self._epoch_to_timestamp_cols: List[str] = []
+        self._epoch_to_timestamp_cols: list[str] = []
         self._add_id_column: Optional[sql_identifier.SqlIdentifier] = None
         self._training_spine_table: str = ""
 
@@ -53,7 +53,7 @@ class ExampleHelper:
                 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]:
+    def load_draft_feature_views(self) -> list[FeatureView]:
         """Return all feature views in an example.
 
         Returns:
@@ -74,7 +74,7 @@ class ExampleHelper:
 
         return fvs
 
-    def load_entities(self) -> List[Entity]:
+    def load_entities(self) -> list[Entity]:
         """Return all entities in an example.
 
         Returns:
@@ -88,7 +88,7 @@ class ExampleHelper:
         with open(file_path) as fs:
             return yaml.safe_load(fs)
 
-    def _create_file_format(self, format_dict: Dict[str, str], format_name: str) -> None:
+    def _create_file_format(self, format_dict: dict[str, str], format_name: str) -> None:
         """Create a file name with given name."""
         self._session.sql(
             f"""
@@ -110,7 +110,7 @@ class ExampleHelper:
             """
         ).collect()
 
-    def _load_csv(self, schema_dict: Dict[str, str], temp_stage_name: str) -> List[str]:
+    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 = ""
@@ -145,7 +145,7 @@ class ExampleHelper:
 
         return [schema_dict["destination_table_name"]]
 
-    def _load_parquet(self, schema_dict: Dict[str, str], temp_stage_name: str) -> List[str]:
+    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"])]
@@ -187,7 +187,7 @@ class ExampleHelper:
 
         return result
 
-    def _load_source_data(self, schema_yaml_file: str) -> List[str]:
+    def _load_source_data(self, schema_yaml_file: str) -> list[str]:
         """Parse a yaml schema file and load data into Snowflake.
 
         Args:
@@ -210,7 +210,7 @@ class ExampleHelper:
         else:
             return self._load_csv(schema_dict, temp_stage_name)
 
-    def load_example(self, example_name: str) -> List[str]:
+    def load_example(self, example_name: str) -> list[str]:
         """Select the active example and load its datasets to Snowflake.
 
         Args:
@@ -247,7 +247,7 @@ class ExampleHelper:
 
         return self.load_source_data(source_yaml_data)
 
-    def load_source_data(self, source_data_name: str) -> List[str]:
+    def load_source_data(self, source_data_name: str) -> list[str]:
         """Load source data into Snowflake.
 
         Args:
@@ -269,10 +269,10 @@ class ExampleHelper:
     def get_current_schema(self) -> str:
         return self._dataset_schema
 
-    def get_label_cols(self) -> List[str]:
+    def get_label_cols(self) -> list[str]:
         return [item.resolved() for item in self._label_columns]
 
-    def get_excluded_cols(self) -> List[str]:
+    def get_excluded_cols(self) -> list[str]:
         return [item.resolved() for item in self._excluded_columns]
 
     def get_training_data_timestamp_col(self) -> Optional[str]:

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

@@ -1,5 +1,3 @@
-from typing import List
-
 from snowflake.ml.feature_store import Entity
 
 trip_id = Entity(name="TRIP_ID", join_keys=["TRIP_ID"], desc="Trip id.")
@@ -8,5 +6,5 @@ location_id = Entity(name="DOLOCATIONID", join_keys=["DOLOCATIONID"], desc="Drop
 
 
 # This will be invoked by example_helper.py. Do not change function name.
-def get_all_entities() -> List[Entity]:
+def get_all_entities() -> list[Entity]:
     return [trip_id, location_id]

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

@@ -1,5 +1,3 @@
-from typing import List
-
 from snowflake.ml.feature_store import FeatureView
 from snowflake.ml.feature_store.examples.new_york_taxi_features.entities import (
     location_id,
@@ -9,7 +7,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a draft feature view."""
     feature_df = session.sql(

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

@@ -1,5 +1,3 @@
-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
@@ -7,7 +5,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a draft feature view."""
     feature_df = session.sql(

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

@@ -1,5 +1,3 @@
-from typing import List
-
 from snowflake.ml.feature_store import Entity
 
 wine_id = Entity(
@@ -10,5 +8,5 @@ wine_id = Entity(
 
 
 # This will be invoked by example_helper.py. Do not change function name.
-def get_all_entities() -> List[Entity]:
+def get_all_entities() -> list[Entity]:
     return [wine_id]

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

@@ -1,5 +1,3 @@
-from typing import List
-
 from snowflake.ml.feature_store import FeatureView
 from snowflake.ml.feature_store.examples.wine_quality_features.entities import wine_id
 from snowflake.snowpark import DataFrame, Session, functions as F
@@ -7,7 +5,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a feature view about trip station."""
     feature_df = source_dfs[0].select(

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

@@ -1,5 +1,3 @@
-from typing import List
-
 from snowflake.ml.feature_store import FeatureView
 from snowflake.ml.feature_store.examples.wine_quality_features.entities import wine_id
 from snowflake.snowpark import DataFrame, Session
@@ -7,7 +5,7 @@ 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], database: str, schema: str
+    session: Session, source_dfs: list[DataFrame], source_tables: list[str], database: str, schema: str
 ) -> FeatureView:
     """Create a feature view about trip station."""
     feature_df = source_dfs[0].select("WINE_ID", "SULPHATES", "ALCOHOL")