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

snowflake/ml/data/_internal/arrow_ingestor.py

@@ -0,0 +1,284 @@
+import collections
+import logging
+import os
+import time
+from typing import Any, Deque, Dict, Iterator, List, Optional, Union
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+import pyarrow as pa
+import pyarrow.dataset as pds
+
+from snowflake import snowpark
+from snowflake.connector import result_batch
+from snowflake.ml.data import data_ingestor, data_source, ingestor_utils
+
+_EMPTY_RECORD_BATCH = pa.RecordBatch.from_arrays([], [])
+
+# The row count for batches read from PyArrow Dataset. This number should be large enough so that
+# dataset.to_batches() would read in a very large portion of, if not entirely, a parquet file.
+_DEFAULT_DATASET_BATCH_SIZE = 1000000
+
+
+class _RecordBatchesBuffer:
+    """A queue that stores record batches and tracks the total num of rows in it."""
+
+    def __init__(self) -> None:
+        self.buffer: Deque[pa.RecordBatch] = collections.deque()
+        self.num_rows = 0
+
+    def append(self, rb: pa.RecordBatch) -> None:
+        self.buffer.append(rb)
+        self.num_rows += rb.num_rows
+
+    def appendleft(self, rb: pa.RecordBatch) -> None:
+        self.buffer.appendleft(rb)
+        self.num_rows += rb.num_rows
+
+    def popleft(self) -> pa.RecordBatch:
+        popped = self.buffer.popleft()
+        self.num_rows -= popped.num_rows
+        return popped
+
+
+class ArrowIngestor(data_ingestor.DataIngestor):
+    """Read and parse the data sources into an Arrow Dataset and yield batched numpy array in dict."""
+
+    def __init__(
+        self,
+        session: snowpark.Session,
+        data_sources: List[data_source.DataSource],
+        format: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Args:
+            session: The Snowpark Session to use.
+            data_sources: List of data sources to ingest.
+            format: Currently “parquet”, “ipc”/”arrow”/”feather”, “csv”, “json”, and “orc” are supported.
+                Will be inferred if not specified.
+            kwargs: Miscellaneous arguments passed to underlying PyArrow Dataset initializer.
+        """
+        self._session = session
+        self._data_sources = data_sources
+        self._format = format
+        self._kwargs = kwargs
+
+        self._schema: Optional[pa.Schema] = None
+
+    @classmethod
+    def from_sources(cls, session: snowpark.Session, sources: List[data_source.DataSource]) -> "ArrowIngestor":
+        return cls(session, sources)
+
+    @property
+    def data_sources(self) -> List[data_source.DataSource]:
+        return self._data_sources
+
+    def to_batches(
+        self,
+        batch_size: int,
+        shuffle: bool = True,
+        drop_last_batch: bool = True,
+    ) -> Iterator[Dict[str, npt.NDArray[Any]]]:
+        """Iterate through PyArrow Dataset to generate batches whose length equals to expected batch size.
+
+        As we are generating batches with the exactly same length, the last few rows in each file might get left as they
+        are not long enough to form a batch. These rows will be put into a temporary buffer and combine with the first
+        few rows of the next file to generate a new batch.
+
+        Args:
+            batch_size: Specifies the size of each batch that will be yield
+            shuffle: Whether the data in the file will be shuffled. If set to be true, it will first randomly shuffle
+                the order of files, and then shuflle the order of rows in each file.
+            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.
+
+        Yields:
+            A dict mapping column names to the corresponding data fetch from that column.
+        """
+        self._rb_buffer = _RecordBatchesBuffer()
+
+        # Extract schema if not already known
+        dataset = self._get_dataset(shuffle)
+        if self._schema is None:
+            self._schema = dataset.schema
+
+        for rb in _retryable_batches(dataset, batch_size=max(_DEFAULT_DATASET_BATCH_SIZE, batch_size)):
+            if shuffle:
+                rb = rb.take(np.random.permutation(rb.num_rows))
+            self._rb_buffer.append(rb)
+            while self._rb_buffer.num_rows >= batch_size:
+                yield self._get_batches_from_buffer(batch_size)
+
+        if self._rb_buffer.num_rows and not drop_last_batch:
+            yield self._get_batches_from_buffer(batch_size)
+
+    def to_pandas(self, limit: Optional[int] = None) -> pd.DataFrame:
+        ds = self._get_dataset(shuffle=False)
+        table = ds.to_table() if limit is None else ds.head(num_rows=limit)
+        return table.to_pandas()
+
+    def _get_dataset(self, shuffle: bool) -> pds.Dataset:
+        format = self._format
+        sources: List[Any] = []
+        source_format = None
+        for source in self._data_sources:
+            if isinstance(source, str):
+                sources.append(source)
+                source_format = format or os.path.splitext(source)[-1]
+            elif isinstance(source, data_source.DatasetInfo):
+                if not self._kwargs.get("filesystem"):
+                    self._kwargs["filesystem"] = ingestor_utils.get_dataset_filesystem(self._session, source)
+                sources.extend(
+                    ingestor_utils.get_dataset_files(self._session, source, filesystem=self._kwargs["filesystem"])
+                )
+                source_format = "parquet"
+            elif isinstance(source, data_source.DataFrameInfo):
+                # FIXME: This currently loads all result batches into memory so that it
+                #        can be passed into pyarrow.dataset as a list/tuple of pa.RecordBatches
+                #        We may be able to optimize this by splitting the result batches into
+                #        in-memory (first batch) and file URLs (subsequent batches) and creating a
+                #        union dataset.
+                result_batches = ingestor_utils.get_dataframe_result_batches(self._session, source)
+                sources.extend(
+                    b.to_arrow(self._session.connection)
+                    if isinstance(b, result_batch.ArrowResultBatch)
+                    else b.to_arrow()
+                    for b in result_batches
+                )
+                # HACK: Mitigate typing inconsistencies in Snowpark results
+                if len(sources) > 0:
+                    sources = [_cast_if_needed(s, sources[-1].schema) for s in sources]
+                source_format = None  # Arrow Dataset expects "None" for in-memory datasets
+            else:
+                raise RuntimeError(f"Unsupported data source type: {type(source)}")
+
+            # Make sure source types not mixed
+            if format and format != source_format:
+                raise RuntimeError(f"Unexpected data source format (expected {format}, found {source_format})")
+            format = source_format
+
+        # Re-shuffle input files on each iteration start
+        if shuffle:
+            np.random.shuffle(sources)
+        pa_dataset: pds.Dataset = pds.dataset(sources, format=format, **self._kwargs)
+        return pa_dataset
+
+    def _get_batches_from_buffer(self, batch_size: int) -> Dict[str, npt.NDArray[Any]]:
+        """Generate new batches from the existing record batch buffer."""
+        cnt_rbs_num_rows = 0
+        candidates = []
+
+        # Keep popping record batches in buffer until there are enough rows for a batch.
+        while self._rb_buffer.num_rows and cnt_rbs_num_rows < batch_size:
+            candidate = self._rb_buffer.popleft()
+            cnt_rbs_num_rows += candidate.num_rows
+            candidates.append(candidate)
+
+        # When there are more rows than needed, slice the last popped batch to fit batch_size.
+        if cnt_rbs_num_rows > batch_size:
+            row_diff = cnt_rbs_num_rows - batch_size
+            slice_target = candidates[-1]
+            cut_off = slice_target.num_rows - row_diff
+            to_merge = slice_target.slice(length=cut_off)
+            left_over = slice_target.slice(offset=cut_off)
+            candidates[-1] = to_merge
+            self._rb_buffer.appendleft(left_over)
+
+        res = _merge_record_batches(candidates)
+        return _record_batch_to_arrays(res)
+
+
+def _merge_record_batches(record_batches: List[pa.RecordBatch]) -> pa.RecordBatch:
+    """Merge a list of arrow RecordBatches into one. Similar to MergeTables."""
+    if not record_batches:
+        return _EMPTY_RECORD_BATCH
+    if len(record_batches) == 1:
+        return record_batches[0]
+    record_batches = list(filter(lambda rb: rb.num_rows > 0, record_batches))
+    one_chunk_table = pa.Table.from_batches(record_batches).combine_chunks()
+    batches = one_chunk_table.to_batches(max_chunksize=None)
+    return batches[0]
+
+
+def _record_batch_to_arrays(rb: pa.RecordBatch) -> Dict[str, npt.NDArray[Any]]:
+    """Transform the record batch to a (string, numpy array) dict."""
+    batch_dict = {}
+    for column, column_schema in zip(rb, rb.schema):
+        # zero_copy_only=False because of nans. Ideally nans should have been imputed in feature engineering.
+        array = column.to_numpy(zero_copy_only=False)
+        batch_dict[column_schema.name] = array
+    return batch_dict
+
+
+def _retryable_batches(
+    dataset: pds.Dataset, batch_size: int, max_retries: int = 3, delay: int = 0
+) -> Iterator[pa.RecordBatch]:
+    """Make the Dataset to_batches retryable."""
+    retries = 0
+    current_batch_index = 0
+
+    while True:
+        try:
+            for batch_index, batch in enumerate(dataset.to_batches(batch_size=batch_size)):
+                if batch_index < current_batch_index:
+                    # Skip batches that have already been processed
+                    continue
+
+                yield batch
+                current_batch_index = batch_index + 1
+            # Exit the loop once all batches are processed
+            break
+
+        except Exception as e:
+            if retries < max_retries:
+                retries += 1
+                logging.info(f"Error encountered: {e}. Retrying {retries}/{max_retries}...")
+                time.sleep(delay)
+            else:
+                raise e
+
+
+def _cast_if_needed(
+    batch: Union[pa.Table, pa.RecordBatch], schema: Optional[pa.Schema] = None
+) -> Union[pa.Table, pa.RecordBatch]:
+    """
+    Cast the batch to be compatible with downstream frameworks. Returns original batch if cast is not necessary.
+    Besides casting types to match `schema` (if provided), this function also applies the following casting:
+        - Decimal (fixed-point) types: Convert to float or integer types based on scale and byte length
+
+    Args:
+        batch: The PyArrow batch to cast if needed
+        schema: Optional schema the batch should be casted to match. Note that compatibility type casting takes
+            precedence over the provided schema, e.g. if the schema has decimal types the result will be further
+            cast into integer/float types.
+
+    Returns:
+        The type-casted PyArrow batch, or the original batch if casting was not necessary
+    """
+    schema = schema or batch.schema
+    assert len(batch.schema) == len(schema)
+    fields = []
+    cast_needed = False
+    for field, target in zip(batch.schema, schema):
+        # Need to convert decimal types to supported types. This behavior supersedes target schema data types
+        if pa.types.is_decimal(target.type):
+            byte_length = int(target.metadata.get(b"byteLength", 8))
+            if int(target.metadata.get(b"scale", 0)) > 0:
+                target = target.with_type(pa.float32() if byte_length == 4 else pa.float64())
+            else:
+                if byte_length == 2:
+                    target = target.with_type(pa.int16())
+                elif byte_length == 4:
+                    target = target.with_type(pa.int32())
+                else:  # Cap out at 64-bit
+                    target = target.with_type(pa.int64())
+        if not field.equals(target):
+            cast_needed = True
+            field = target
+        fields.append(field)
+
+    if cast_needed:
+        return batch.cast(pa.schema(fields))
+    return batch

snowflake/ml/data/data_connector.py

@@ -0,0 +1,186 @@
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Type, TypeVar
+
+import numpy.typing as npt
+from typing_extensions import deprecated
+
+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
+
+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"""
+
+    DEFAULT_INGESTOR_CLASS: Type[data_ingestor.DataIngestor] = ArrowIngestor
+
+    def __init__(
+        self,
+        ingestor: data_ingestor.DataIngestor,
+    ) -> None:
+        self._ingestor = ingestor
+
+    @classmethod
+    @snowpark._internal.utils.private_preview(version="1.6.0")
+    def from_dataframe(
+        cls: Type[DataConnectorType],
+        df: snowpark.DataFrame,
+        ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None,
+        **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
+        return cls.from_sources(df._session, [source], ingestor_class=ingestor_class, **kwargs)
+
+    @classmethod
+    def from_dataset(
+        cls: Type[DataConnectorType],
+        ds: "dataset.Dataset",
+        ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None,
+        **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)
+        )
+        return cls.from_sources(ds._session, [source], ingestor_class=ingestor_class, **kwargs)
+
+    @classmethod
+    @telemetry.send_api_usage_telemetry(
+        project=_PROJECT,
+        subproject_extractor=lambda cls: cls.__name__,
+        func_params_to_log=["sources", "ingestor_class"],
+    )
+    def from_sources(
+        cls: Type[DataConnectorType],
+        session: snowpark.Session,
+        sources: List[data_source.DataSource],
+        ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None,
+        **kwargs: Any
+    ) -> DataConnectorType:
+        ingestor_class = ingestor_class or cls.DEFAULT_INGESTOR_CLASS
+        ingestor = ingestor_class.from_sources(session, sources)
+        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)
+
+    @deprecated(
+        "to_torch_datapipe() is deprecated and will be removed in a future release. Use to_torch_dataset() instead"
+    )
+    @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=["shuffle"],
+    )
+    def to_torch_dataset(self, *, shuffle: bool = False) -> "torch_data.IterableDataset":  # type: ignore[type-arg]
+        """Transform the Snowflake data into a PyTorch Iterable Dataset to be used with a DataLoader.
+
+        Return a PyTorch Dataset which iterates on rows of data.
+
+        Args:
+            shuffle: It specifies whether the data will be shuffled. If True, files will be shuffled, and
+                rows in each file will also be shuffled.
+
+        Returns:
+            A PyTorch Iterable Dataset that yields data.
+        """
+        from snowflake.ml.data import torch_dataset
+
+        return torch_dataset.TorchDataset(self._ingestor, shuffle)
+
+    @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,45 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Protocol,
+    Type,
+    TypeVar,
+)
+
+from numpy import typing as npt
+
+from snowflake import snowpark
+from snowflake.ml.data import data_source
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+DataIngestorType = TypeVar("DataIngestorType", bound="DataIngestor")
+
+
+class DataIngestor(Protocol):
+    @classmethod
+    def from_sources(
+        cls: Type[DataIngestorType], session: snowpark.Session, sources: List[data_source.DataSource]
+    ) -> DataIngestorType:
+        raise NotImplementedError
+
+    @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/data/ingestor_utils.py

@@ -0,0 +1,62 @@
+from typing import List, Optional
+
+import fsspec
+
+from snowflake import snowpark
+from snowflake.connector import result_batch
+from snowflake.ml.data import data_source
+from snowflake.ml.fileset import snowfs
+
+_TARGET_FILE_SIZE = 32 * 2**20  # The max file size for data loading.
+
+
+def get_dataframe_result_batches(
+    session: snowpark.Session, df_info: data_source.DataFrameInfo
+) -> List[result_batch.ResultBatch]:
+    """Retrieve the ResultBatches for a given query"""
+    cursor = session._conn._cursor
+
+    if df_info.query_id:
+        query_id = df_info.query_id
+    else:
+        query_id = session.sql(df_info.sql).collect_nowait().query_id
+
+    # TODO: Check if query result cache is still live
+    cursor.get_results_from_sfqid(sfqid=query_id)
+
+    # Prefetch hook should be set by `get_results_from_sfqid`
+    # This call blocks until the query results are ready
+    if cursor._prefetch_hook is None:
+        raise RuntimeError("Loading data from result query failed unexpectedly. Please contact Snowflake support.")
+    cursor._prefetch_hook()
+    batches = cursor.get_result_batches()
+    if batches is None:
+        raise ValueError(
+            "Failed to retrieve training data. Query status:" f" {session._conn._conn.get_query_status(query_id)}"
+        )
+    return batches
+
+
+def get_dataset_filesystem(
+    session: snowpark.Session, ds_info: Optional[data_source.DatasetInfo] = None
+) -> fsspec.AbstractFileSystem:
+    """Get the fsspec filesystem for a given Dataset"""
+    # We can't directly load the Dataset to avoid a circular dependency
+    # Dataset -> DatasetReader -> DataConnector -> DataIngestor -> (?) ingestor_utils -> Dataset
+    # TODO: Automatically pick appropriate fsspec implementation based on protocol in URL
+    return snowfs.SnowFileSystem(
+        snowpark_session=session,
+        cache_type="bytes",
+        block_size=2 * _TARGET_FILE_SIZE,
+    )
+
+
+def get_dataset_files(
+    session: snowpark.Session, ds_info: data_source.DatasetInfo, filesystem: Optional[fsspec.AbstractFileSystem] = None
+) -> List[str]:
+    """Get the list of files in a given Dataset"""
+    if filesystem is None:
+        filesystem = get_dataset_filesystem(session, ds_info)
+    assert bool(ds_info.url)  # Not null or empty
+    files = sorted(filesystem.ls(ds_info.url))
+    return [filesystem.unstrip_protocol(f) for f in files]

snowflake/ml/data/torch_dataset.py

@@ -0,0 +1,33 @@
+from typing import Any, Dict, Iterator
+
+import torch.utils.data
+
+from snowflake.ml.data import data_ingestor
+
+
+class TorchDataset(torch.utils.data.IterableDataset[Dict[str, Any]]):
+    """Implementation of PyTorch IterableDataset"""
+
+    def __init__(self, ingestor: data_ingestor.DataIngestor, shuffle: bool = False) -> None:
+        """Not intended for direct usage. Use DataConnector.to_torch_dataset() instead"""
+        self._ingestor = ingestor
+        self._shuffle = shuffle
+
+    def __iter__(self) -> Iterator[Dict[str, Any]]:
+        max_idx = 0
+        filter_idx = 0
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is not None:
+            max_idx = worker_info.num_workers - 1
+            filter_idx = worker_info.id
+
+        counter = 0
+        for batch in self._ingestor.to_batches(batch_size=1, shuffle=self._shuffle, drop_last_batch=False):
+            # Skip indices during multi-process data loading to prevent data duplication
+            if counter == filter_idx:
+                yield {k: v.item() for k, v in batch.items()}
+
+            if counter < max_idx:
+                counter += 1
+            else:
+                counter = 0

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,
@@ -177,18 +176,7 @@ class Dataset(lineage_node.LineageNode):
                 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._lineage_node_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

snowflake/ml/dataset/dataset_metadata.py

@@ -15,11 +15,13 @@ class FeatureStoreMetadata:
     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.
+        compact_feature_views: A compact representation of a FeatureView or FeatureViewSlice.
         spine_timestamp_col: Timestamp column which was used for point-in-time correct feature lookup.
     """
 
     spine_query: str
-    serialized_feature_views: List[str]
+    serialized_feature_views: Optional[List[str]] = None
+    compact_feature_views: Optional[List[str]] = None
     spine_timestamp_col: Optional[str] = None
 
     def to_json(self) -> str: