snowflake-ml-python 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

snowflake/ml/dataset/dataset_metadata.py

@@ -0,0 +1,103 @@
+import dataclasses
+import json
+import typing
+from typing import Any, Dict, List, Optional, Union
+
+_PROPERTY_TYPE_KEY = "$proptype$"
+DATASET_SCHEMA_VERSION = "1"
+
+
+@dataclasses.dataclass(frozen=True)
+class FeatureStoreMetadata:
+    """
+    Feature store metadata.
+
+    Properties:
+        spine_query: The input query on source table which will be joined with features.
+        serialized_feature_views: A list of serialized feature objects in the feature store.
+        spine_timestamp_col: Timestamp column which was used for point-in-time correct feature lookup.
+    """
+
+    spine_query: str
+    serialized_feature_views: List[str]
+    spine_timestamp_col: Optional[str] = None
+
+    def to_json(self) -> str:
+        return json.dumps(dataclasses.asdict(self))
+
+    @classmethod
+    def from_json(cls, input_json: Union[Dict[str, Any], str, bytes]) -> "FeatureStoreMetadata":
+        if isinstance(input_json, dict):
+            return cls(**input_json)
+        return cls(**json.loads(input_json))
+
+
+DatasetPropertiesType = Union[
+    FeatureStoreMetadata,
+]
+
+# Union[T] gets automatically squashed to T, so default to [T] if get_args() returns empty
+_DatasetPropTypes = typing.get_args(DatasetPropertiesType) or [DatasetPropertiesType]
+_DatasetPropTypeDict = {t.__name__: t for t in _DatasetPropTypes}
+
+
+@dataclasses.dataclass(frozen=True)
+class DatasetMetadata:
+    """
+    Dataset metadata.
+
+    Properties:
+        source_query: The query string used to produce the Dataset.
+        owner: The owner of the Dataset.
+        generation_timestamp: The timestamp when this dataset was generated.
+        exclude_cols: Name of column(s) in dataset to be excluded during training/testing.
+            These are typically columns for human inspection such as timestamp or other meta-information.
+            Columns included in `label_cols` do not need to be included here.
+        label_cols: Name of column(s) in dataset that contains labels.
+        properties: Additional metadata properties.
+    """
+
+    source_query: str
+    owner: str
+    exclude_cols: Optional[List[str]] = None
+    label_cols: Optional[List[str]] = None
+    properties: Optional[DatasetPropertiesType] = None
+    schema_version: str = dataclasses.field(default=DATASET_SCHEMA_VERSION, init=False)
+
+    def to_json(self) -> str:
+        state_dict = dataclasses.asdict(self)
+        if self.properties:
+            prop_type = type(self.properties).__name__
+            if prop_type not in _DatasetPropTypeDict:
+                raise ValueError(
+                    f"Unsupported `properties` type={prop_type} (supported={','.join(_DatasetPropTypeDict.keys())})"
+                )
+            state_dict[_PROPERTY_TYPE_KEY] = prop_type
+        return json.dumps(state_dict)
+
+    @classmethod
+    def from_json(cls, input_json: Union[Dict[str, Any], str, bytes]) -> "DatasetMetadata":
+        if not input_json:
+            raise ValueError("json_str was empty or None")
+        try:
+            state_dict: Dict[str, Any] = (
+                input_json if isinstance(input_json, dict) else json.loads(input_json, strict=False)
+            )
+
+            # TODO: Validate schema version
+            _ = state_dict.pop("schema_version", DATASET_SCHEMA_VERSION)
+
+            prop_type = state_dict.pop(_PROPERTY_TYPE_KEY, None)
+            prop_values = state_dict.get("properties", {})
+            if prop_type:
+                prop_cls = _DatasetPropTypeDict.get(prop_type, None)
+                if prop_cls is None:
+                    raise TypeError(
+                        f"Unsupported `properties` type={prop_type} (supported={','.join(_DatasetPropTypeDict.keys())})"
+                    )
+                state_dict["properties"] = prop_cls(**prop_values)
+            elif prop_values:
+                raise TypeError(f"`properties` provided but missing `{_PROPERTY_TYPE_KEY}`")
+            return cls(**state_dict)
+        except TypeError as e:
+            raise ValueError("Invalid input schema") from e

snowflake/ml/dataset/dataset_reader.py

@@ -0,0 +1,202 @@
+from typing import Any, List
+
+import pandas as pd
+
+from snowflake import snowpark
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.lineage import data_source, dataset_dataframe
+from snowflake.ml._internal.utils import import_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:
+    """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],
+    ) -> None:
+        """Initialize a DatasetVersion object.
+
+        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] = []
+
+    def _list_files(self) -> List[str]:
+        """Private helper function that lists all files in this DatasetVersion and caches the results."""
+        if self._files:
+            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]
+        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.
+
+        The file paths follows the snow protocol.
+
+        Returns:
+            A list of remote file paths
+
+        Example:
+        >>> dsv.files()
+        ----
+        ["snow://dataset/mydb.myschema.mydataset/versions/test/data_0_0_0.snappy.parquet",
+         "snow://dataset/mydb.myschema.mydataset/versions/test/data_0_0_1.snappy.parquet"]
+        """
+        files = self._list_files()
+        return [self._fs.unstrip_protocol(f) for f in files]
+
+    @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
+    def filesystem(self) -> snowfs.SnowFileSystem:
+        """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,
+        func_params_to_log=["only_feature_cols"],
+    )
+    def to_snowpark_dataframe(self, only_feature_cols: bool = False) -> snowpark.DataFrame:
+        """Convert the DatasetVersion to a Snowpark DataFrame.
+
+        Args:
+            only_feature_cols: If True, drops exclude_cols and label_cols from returned DataFrame.
+                The original DatasetVersion is unaffected.
+
+        Returns:
+            A Snowpark dataframe that contains the data of this DatasetVersion.
+
+        Note: The dataframe generated by this method might not have the same schema as the original one. Specifically,
+            - NUMBER type with scale != 0 will become float.
+            - Unsupported types (see comments of :func:`Dataset.create_version`) will not have any guarantee.
+                For example, an OBJECT column may be scanned back as a STRING column.
+        """
+        file_path_pattern = ".*data_.*[.]parquet"
+        dfs: List[snowpark.DataFrame] = []
+        for source in self._sources:
+            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)
+            dfs.append(df)
+
+        combined_df = dfs[0]
+        for df in dfs[1:]:
+            combined_df = combined_df.union_all_by_name(df)
+        return dataset_dataframe.DatasetDataFrame.from_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_dfs = []
+        for file in files:
+            with self._fs.open(file) as fp:
+                pd_dfs.append(pd.read_parquet(fp))
+        pd_df = pd_dfs[0] if len(pd_dfs) == 1 else pd.concat(pd_dfs, ignore_index=True, copy=False)
+        return pd_df