dataframe_to_mds fails on large datasets due to Worker Disk Exhaustion

Issue Description

Calling dataframe_to_mds on a large Spark DataFrame leads to worker disk exhaustion: no space left of disk. This happens when a remote path is provided in the out argument of mds_kwargs.

Following is the code snippet I had used:

train_df = spark.read.parquet(train_data_path)

# remote output path
train_shard_path = "gs://<bucket_name><prefix>/train"

mds_kwargs = {
    'out': train_shard_path,
    'columns': {
        'X': 'int64',
        'Y': 'ndarray:float32',
        'A': 'int64',
        'B': 'ndarray:float32',
        'C': 'ndarray:int64',
        'D': 'ndarray:float32'
    },
    'compression': 'zstd',
    'hashes': ['sha1'],
    'keep_local': False 
}


mds_path = dataframe_to_mds(
    train_df,
    merge_index=True,
    mds_kwargs=mds_kwargs
)

The current implementation is dependent on the worker's disk space (or shared volume size). If the worker disk space is less than the total size of the dataset (dataframe) then the dataframe_to_mds method crashes as the workers run out of disk space.

Environment

• OS: [Ubuntu 20.04]
• Spark version: 3.5

Expected behavior

The dataframe_to_mds function should be able to process datasets of any size, limited only by remote storage capacity rather than local worker disk or driver memory.

Start at the dataframe_to_mds entry point and reproduce the failure with a large Spark DataFrame, a remote out path, merge_index=True, and keep_local=False. Trace how worker disk and remote output are used; done means processing is no longer limited by local worker disk or driver memory, subject to remote storage capacity.