.. _data-curator-sparkother:

==========================================================================
Reading and writing datasets with NeMo Curator and Apache Spark
==========================================================================
-----------------------------------------
Background
-----------------------------------------

NeMo Curator uses the ``DocumentDataset`` class to read and write JSONL and Parquet files.
It is a wrapper around a `Dask (or Dask-cuDF) DataFrame <https://docs.dask.org/en/stable/dataframe.html>`_.  Apache Spark can read and write JSONL and Parquet files generated by NeMo Curator, and similarly, NeMo Curator can work with the outputs generated by Spark.

-----------------------------------------
Usage
-----------------------------------------

To demonstrate how this would work, consider the following example:

.. code-block:: python

    import dask.dataframe as dd
    import pandas as pd
    from nemo_curator.datasets import DocumentDataset

    # Create sample data
    data = {
        "id": [1, 2, 3],
        "text": [
            "This is a tiny story.",
            "Another tiny story appears here.",
            "Yet another tiny story for you."
        ]
    }

    # Convert to a pandas DataFrame first
    df = pd.DataFrame(data)

    # Convert pandas DataFrame to DocumentDataset
    stories_ds = DocumentDataset(dd.from_pandas(df, npartitions=2))

    # Write the dataset to JSONL files
    stories_ds.to_json("tiny_stories/", write_to_filename=False)

This will create two JSONL files in the directory ``tiny_stories/``:

::

    tiny_stories/
        0.part
        1.part

Apache Spark can read these files using standard APIs. Let's first create a Spark session called ``NeMoCuratorExample``, then we can read files in the directory using:

.. code-block:: python

    from pyspark.sql import SparkSession

    spark = SparkSession.builder.appName("NeMoCuratorExample").getOrCreate()

    # Reading JSONL file
    stories_df = spark.read.json("tiny_stories")
    stories_df.show()

Let's go ahead and add a couple of columns to the Spark DataFrame:

.. code-block:: python

    from pyspark.sql.functions import size, split, length

    # Calculate Word Count
    stories_df = stories_df.withColumn("WordCount", size(split(stories_df["text"], r"\s+")))

    # Calculate Character Count
    stories_df = stories_df.withColumn("CharacterCount", length(stories_df["text"]))

    stories_df.write.mode("overwrite").parquet("tiny_stories_transformed")

To interoperate between NeMo Curator ``DocumentDataset`` and Spark DataFrames, we recommend using Parquet files for data exchange.
The following code snippet demonstrates how to read output from a Spark DataFrame into a NeMo Curator ``DocumentDataset``:

.. code-block:: python

    from nemo_curator.utils.file_utils import get_all_files_paths_under

    # Ignores checksum and marker files created by Spark
    processed_files = [
         filename for filename in get_all_files_paths_under("tiny_stories_transformed")
         if not (filename.endswith(".crc") or filename.endswith("_SUCCESS"))
    ]

    stories_dataset = DocumentDataset.read_parquet(processed_files, backend="pandas")

It is worth noting that Spark typically tends to create checksum and other marker files which can vary by Spark distribution,
so it is advisable to ignore them when reading data into a NeMo Curator ``DocumentDataset``.