Function make_spark_converter

petastorm/spark/spark_dataset_converter.py:664–736

Convert a spark dataframe into a :class:`SparkDatasetConverter` object. It will materialize a spark dataframe to the directory specified by spark conf 'petastorm.spark.converter.parentCacheDirUrl'. The dataframe will be materialized in parquet format, and we can specify `parquet_row_

(
        df,
        parquet_row_group_size_bytes=DEFAULT_ROW_GROUP_SIZE_BYTES,
        compression_codec=None,
        dtype='float32'
)

def make_spark_converter(
        df,
        parquet_row_group_size_bytes=DEFAULT_ROW_GROUP_SIZE_BYTES,
        compression_codec=None,
        dtype='float32'
):
    """Convert a spark dataframe into a :class:`SparkDatasetConverter` object.
    It will materialize a spark dataframe to the directory specified by
    spark conf 'petastorm.spark.converter.parentCacheDirUrl'.
    The dataframe will be materialized in parquet format, and we can specify
    `parquet_row_group_size_bytes` and `compression_codec` for the parquet
    format. See params documentation for details.

    The returned `SparkDatasetConverter` object will hold the materialized
    dataframe, and can be used to make one or more tensorflow datasets or
    torch dataloaders.

    We can explicitly delete the materialized dataframe data, see
    `SparkDatasetConverter.delete`, and when the spark application exit,
    it will try best effort to delete the materialized dataframe data.

    :param df: The :class:`pyspark.sql.DataFrame` object to be converted,
               or a string of path pointing to the directory that stores the dataframe data
               as parquet format, on databricks runtime, the path must be a dbfs
               fuse path like 'file:/dbfs/xxx' or a dbfs path like 'dbfs:/xxx'.
    :param parquet_row_group_size_bytes: An int denoting the number of bytes
        in a parquet row group when materializing the dataframe.
    :param compression_codec: Specify compression codec.
        It can be one of 'uncompressed', 'bzip2', 'gzip', 'lz4', 'snappy', 'deflate'.
        Default ``None``. If ``None``, it will leave the data uncompressed.
    :param dtype: ``None``, 'float32' or 'float64', specifying the precision of the floating-point
        elements in the output dataset. Integer types will remain unchanged. If ``None``, all types
        will remain unchanged. Default 'float32'.

    :return: a :class:`SparkDatasetConverter` object that holds the
        materialized dataframe and can be used to make one or more tensorflow
        datasets or torch dataloaders.
    """

    parent_cache_dir_url = _get_parent_cache_dir_url()

    if isinstance(df, str):
        dataset_dir_url = df
        if 'DATABRICKS_RUNTIME_VERSION' in os.environ:
            dataset_dir_url = _normalize_databricks_dbfs_url(
                dataset_dir_url,
                "On databricks runtime, if `df` argument is a string, it must be a dbfs "
                "fuse path like 'file:/dbfs/xxx' or a dbfs path like 'dbfs:/xxx'."
            )
    else:
        # TODO: Improve default behavior to be automatically choosing the best way.
        compression_codec = compression_codec or "uncompressed"

        if compression_codec.lower() not in \
                ['uncompressed', 'bzip2', 'gzip', 'lz4', 'snappy', 'deflate']:
            raise RuntimeError(
                "compression_codec should be None or one of the following values: "
                "'uncompressed', 'bzip2', 'gzip', 'lz4', 'snappy', 'deflate'")