mlflow.data
============

The ``mlflow.data`` module helps you record your model training and evaluation datasets to
runs with MLflow Tracking, as well as retrieve dataset information from runs. It provides the
following important interfaces:

* :py:class:`Dataset <mlflow.data.dataset.Dataset>`: Represents a dataset used in model training or
  evaluation, including features, targets, predictions, and metadata such as the dataset's name, digest (hash)
  schema, profile, and source. You can log this metadata to a run in MLflow Tracking using
  the :py:func:`mlflow.log_input()` API. ``mlflow.data`` provides APIs for constructing
  :py:class:`Datasets <mlflow.data.dataset.Dataset>` from a variety of Python data objects, including
  Pandas DataFrames (:py:func:`mlflow.data.from_pandas()`), NumPy arrays
  (:py:func:`mlflow.data.from_numpy()`), Spark DataFrames (:py:func:`mlflow.data.from_spark()`
  / :py:func:`mlflow.data.load_delta()`), Polars DataFrames (:py:func:`mlflow.data.from_polars()`), and more.

* :py:func:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`: Represents the source of a
  dataset. For example, this may be a directory of files stored in S3, a Delta Table, or a web URL.
  Each :py:class:`Dataset <mlflow.data.dataset.Dataset>` references the source from which it was
  derived. A :py:class:`Dataset <mlflow.data.dataset.Dataset>`'s features and targets may differ
  from the source if transformations and filtering were applied. You can get the
  :py:func:`DatasetSource <mlflow.data.dataset_source.DatasetSource>` of a dataset logged to a
  run in MLflow Tracking using the :py:func:`mlflow.data.get_source()` API.

The following example demonstrates how to use ``mlflow.data`` to log a training dataset to a run,
retrieve information about the dataset from the run, and load the dataset's source.

.. code-block:: python

    import mlflow.data
    import pandas as pd
    from mlflow.data.pandas_dataset import PandasDataset

    # Construct a Pandas DataFrame using iris flower data from a web URL
    dataset_source_url = "http://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
    df = pd.read_csv(dataset_source_url)
    # Construct an MLflow PandasDataset from the Pandas DataFrame, and specify the web URL
    # as the source
    dataset: PandasDataset = mlflow.data.from_pandas(df, source=dataset_source_url)

    with mlflow.start_run():
        # Log the dataset to the MLflow Run. Specify the "training" context to indicate that the
        # dataset is used for model training
        mlflow.log_input(dataset, context="training")

    # Retrieve the run, including dataset information
    run = mlflow.get_run(mlflow.last_active_run().info.run_id)
    dataset_info = run.inputs.dataset_inputs[0].dataset
    print(f"Dataset name: {dataset_info.name}")
    print(f"Dataset digest: {dataset_info.digest}")
    print(f"Dataset profile: {dataset_info.profile}")
    print(f"Dataset schema: {dataset_info.schema}")

    # Load the dataset's source, which downloads the content from the source URL to the local
    # filesystem
    dataset_source = mlflow.data.get_source(dataset_info)
    dataset_source.load()

.. autoclass:: mlflow.data.dataset.Dataset
    :members:
    :undoc-members:
    :show-inheritance:

.. autoclass:: mlflow.data.dataset_source.DatasetSource
    :members:
    :undoc-members:
    :show-inheritance:
    :exclude-members: from_json

    .. method:: from_json(cls, source_json: str) -> DatasetSource

.. autofunction:: mlflow.data.get_source


pandas
~~~~~~

.. autofunction:: mlflow.data.from_pandas

.. autoclass:: mlflow.data.pandas_dataset.PandasDataset()
    :members:
    :undoc-members:
    :exclude-members: to_pyfunc, to_evaluation_dataset


NumPy
~~~~~

.. autofunction:: mlflow.data.from_numpy

.. autoclass:: mlflow.data.numpy_dataset.NumpyDataset()
    :members:
    :undoc-members:
    :exclude-members: to_pyfunc, to_evaluation_dataset


Spark
~~~~~

.. autofunction:: mlflow.data.load_delta

.. autofunction:: mlflow.data.from_spark

.. autoclass:: mlflow.data.spark_dataset.SparkDataset()
    :members:
    :undoc-members:
    :exclude-members: to_pyfunc, to_evaluation_dataset


Hugging Face 
~~~~~~~~~~~~

.. autofunction:: mlflow.data.huggingface_dataset.from_huggingface

.. autoclass:: mlflow.data.huggingface_dataset.HuggingFaceDataset()
    :members:
    :undoc-members:
    :exclude-members: to_pyfunc


TensorFlow 
~~~~~~~~~~~~

.. autofunction:: mlflow.data.tensorflow_dataset.from_tensorflow

.. autoclass:: mlflow.data.tensorflow_dataset.TensorFlowDataset()
    :members:
    :undoc-members:
    :exclude-members: to_pyfunc, 

.. autoclass:: mlflow.data.evaluation_dataset.EvaluationDataset()
    :members:
    :undoc-members:


polars
~~~~~~

.. autofunction:: mlflow.data.from_polars

.. autoclass:: mlflow.data.polars_dataset.PolarsDataset()
    :members:
    :undoc-members:
    :exclude-members: to_pyfunc, to_evaluation_dataset


Dataset Sources 
~~~~~~~~~~~~~~~~

.. autoclass:: mlflow.data.filesystem_dataset_source.FileSystemDatasetSource()
    :members:
    :undoc-members:

.. autoclass:: mlflow.data.http_dataset_source.HTTPDatasetSource()
    :members:
    :undoc-members:
    
.. autoclass:: mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource()
    :members:
    :undoc-members:
    :exclude-members:

.. autoclass:: mlflow.data.delta_dataset_source.DeltaDatasetSource()
    :members:
    :undoc-members:

.. autoclass:: mlflow.data.spark_dataset_source.SparkDatasetSource()
    :members:
    :undoc-members: