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

snowflake/ml/dataset/dataset_reader.py

@@ -1,48 +1,37 @@
-from typing import Any, List
-
-import pandas as pd
-from pyarrow import parquet as pq
+from typing import Any, List, Optional, Type
 
 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, 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.
+        self._session: snowpark.Session = snowpark_session
+        self._fs: snowfs.SnowFileSystem = ingestor_utils.get_dataset_filesystem(self._session)
+        self._files: Optional[List[str]] = None
 
-        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] = []
+    @classmethod
+    def from_dataframe(
+        cls, df: snowpark.DataFrame, ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None, **kwargs: Any
+    ) -> "DatasetReader":
+        # Block superclass constructor from Snowpark DataFrames
+        raise RuntimeError("Creating DatasetReader from DataFrames not supported")
 
     def _list_files(self) -> List[str]:
         """Private helper function that lists all files in this DatasetVersion and caches the results."""
@@ -50,18 +39,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 +70,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 +92,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 +102,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/airline_features/entities.py

@@ -0,0 +1,16 @@
+from typing import List
+
+from snowflake.ml.feature_store import Entity
+
+zipcode_entity = Entity(
+    name="AIRPORT_ZIP_CODE",
+    join_keys=["AIRPORT_ZIP_CODE"],
+    desc="Zip code of the airport.",
+)
+
+plane_entity = Entity(name="PLANE_MODEL", join_keys=["PLANE_MODEL"], desc="The model of an airplane.")
+
+
+# This will be invoked by example_helper.py. Do not change function name.
+def get_all_entities() -> List[Entity]:
+    return [zipcode_entity, plane_entity]

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/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,37 @@
+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,
+            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",  # 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

@@ -0,0 +1,30 @@
+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",
+    ).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

@@ -0,0 +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