Title: Got an error  _pickle.PicklingError use Dataset.from_spark.
State: closed

              python 3.9.2
Got an error  _pickle.PicklingError use Dataset.from_spark.

Did the dataset import load data from spark dataframe using multi-node Spark cluster
df = spark.read.parquet(args.input_data).repartition(50)
ds = Dataset.from_spark(df, keep_in_memory=True,
                        cache_dir="/pnc-data/data/nuplan/t5_spark/cache_data")
ds.save_to_disk(args.output_data)

Error : 
_pickle.PicklingError: Could not serialize object: RuntimeError: It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transforma
tion. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.
23/06/16 21:17:20 WARN ExecutorPodsWatchSnapshotSource: Kubernetes client has been closed (this is expected if the application is shutting down.)

_Originally posted by @yanzia12138 in https://github.com/huggingface/datasets/issues/5701#issuecomment-1594674306_
 
W
Traceback (most recent call last):
  File "/home/work/main.py", line 100, in <module>
    run(args)
  File "/home/work/main.py", line 80, in run
    ds = Dataset.from_spark(df1, keep_in_memory=True,
  File "/home/work/.local/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 1281, in from_spark
    return SparkDatasetReader(
  File "/home/work/.local/lib/python3.9/site-packages/datasets/io/spark.py", line 53, in read
    self.builder.download_and_prepare(
  File "/home/work/.local/lib/python3.9/site-packages/datasets/builder.py", line 909, in download_and_prepare
    self._download_and_prepare(
  File "/home/work/.local/lib/python3.9/site-packages/datasets/builder.py", line 1004, in _download_and_prepare
    self._prepare_split(split_generator, **prepare_split_kwargs)
  File "/home/work/.local/lib/python3.9/site-packages/datasets/packaged_modules/spark/spark.py", line 254, in _prepare_split
    self._validate_cache_dir()
  File "/home/work/.local/lib/python3.9/site-packages/datasets/packaged_modules/spark/spark.py", line 122, in _validate_cache_dir
    self._spark.sparkContext.parallelize(range(1), 1).mapPartitions(create_cache_and_write_probe).collect()
  File "/home/work/.local/lib/python3.9/site-packages/pyspark/rdd.py", line 950, in collect
    sock_info = self.ctx._jvm.PythonRDD.collectAndServe(self._jrdd.rdd())
  File "/home/work/.local/lib/python3.9/site-packages/pyspark/rdd.py", line 2951, in _jrdd
    wrapped_func = _wrap_function(self.ctx, self.func, self._prev_jrdd_deserializer,
  File "/home/work/.local/lib/python3.9/site-packages/pyspark/rdd.py", line 2830, in _wrap_function
    pickled_command, broadcast_vars, env, includes = _prepare_for_python_RDD(sc, command)
  File "/home/work/.local/lib/python3.9/site-packages/pyspark/rdd.py", line 2816, in _prepare_for_python_RDD
    pickled_command = ser.dumps(command)
  File "/home/work/.local/lib/python3.9/site-packages/pyspark/serializers.py", line 447, in dumps
    raise pickle.PicklingError(msg)
_pickle.PicklingError: Could not serialize object: RuntimeError: It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. S
parkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.
23/06/19 13:51:21 WARN ExecutorPodsWatchSnapshotSource: Kubernetes client has been closed (this is expected if the application is shutting down.)


Comments:
i got error using method from_spark  when using multi-node Spark cluster. seems could only use "from_spark" in local?

@lhoestq 

cc @maddiedawson it looks like there an issue with `_validate_cache_dir` ?

It looks like the function passed to mapPartitions has a reference to the Spark dataset builder, and therefore contains the SparkContext itself.

I think it can be fixed by defining `create_cache_and_write_probe` outside the Spark dataset builder, and pass a `partial(create_cache_and_write_probe, cache_dir=self._cache_dir)` to `mapPartitions`

Just saw this; thanks for flagging! Your proposed solution sounds good. I can prepare a PR

@maddiedawson  can you show me the demo ,so i can test in local .before your PR