GitHub issue: `IterableDataset.from_spark` does not support multiple workers in pytorch `Dataloader`
Author: NazyS (NONE)
State: open | Created: 2023-10-26 12:43:36

### Describe the bug

Looks like `IterableDataset.from_spark` does not support multiple workers in pytorch `Dataloader` if I'm not missing anything.

Also, returns not consistent error messages, which probably depend on the nondeterministic order of worker executions

Some exampes I've encountered:

```
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 79, in __iter__
    yield from self.generate_examples_fn()
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 49, in generate_fn
    df_with_partition_id = df.select("*", pyspark.sql.functions.spark_partition_id().alias("part_id"))
  File "/databricks/spark/python/pyspark/instrumentation_utils.py", line 54, in wrapper
    logger.log_failure(
  File "/databricks/spark/python/pyspark/databricks/usage_logger.py", line 70, in log_failure
    self.logger.recordFunctionCallFailureEvent(
  File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/java_gateway.py", line 1322, in __call__
    return_value = get_return_value(
  File "/databricks/spark/python/pyspark/errors/exceptions/captured.py", line 188, in deco
    return f(*a, **kw)
  File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/protocol.py", line 342, in get_return_value
    return OUTPUT_CONVERTER[type](answer[2:], gateway_client)
KeyError: 'c'
```

```
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 79, in __iter__
    yield from self.generate_examples_fn()
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 49, in generate_fn
    df_with_partition_id = df.select("*", pyspark.sql.functions.spark_partition_id().alias("part_id"))
  File "/databricks/spark/python/pyspark/sql/utils.py", line 162, in wrapped
    return f(*args, **kwargs)
  File "/databricks/spark/python/pyspark/sql/functions.py", line 4893, in spark_partition_id
    return _invoke_function("spark_partition_id")
  File "/databricks/spark/python/pyspark/sql/functions.py", line 98, in _invoke_function
    return Column(jf(*args))
  File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/java_gateway.py", line 1322, in __call__
    return_value = get_return_value(
  File "/databricks/spark/python/pyspark/errors/exceptions/captured.py", line 188, in deco
    return f(*a, **kw)
  File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/protocol.py", line 342, in get_return_value
    return OUTPUT_CONVERTER[type](answer[2:], gateway_client)
KeyError: 'm'
```

```
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 79, in __iter__
    yield from self.generate_examples_fn()
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 49, in generate_fn
    df_with_partition_id = df.select("*", pyspark.sql.functions.spark_partition_id().alias("part_id"))
  File "/databricks/spark/python/pyspark/sql/utils.py", line 162, in wrapped
    return f(*args, **kwargs)
  File "/databricks/spark/python/pyspark/sql/functions.py", line 4893, in spark_partition_id
    return _invoke_function("spark_partition_id")
  File "/databricks/spark/python/pyspark/sql/functions.py", line 97, in _invoke_function
    jf = _get_jvm_function(name, SparkContext._active_spark_context)
  File "/databricks/spark/python/pyspark/sql/functions.py", line 88, in _get_jvm_function
    return getattr(sc._jvm.functions, name)
  File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/java_gateway.py", line 1725, in __getattr__
    raise Py4JError(message)
py4j.protocol.Py4JError: functions does not exist in the JVM
```

### Steps to reproduce the bug

```python
import pandas as pd
import numpy as np

batch_size = 16

pdf = pd.DataFrame({
    key: np.random.rand(16*100) for key in ['feature', 'target']
})
test_df = spark.createDataFrame(pdf)

from datasets import IterableDataset
from torch.utils.data import DataLoader

ids = IterableDataset.from_spark(test_df)
for batch in DataLoader(ids, batch_size=16, num_workers=4):
    for k, b in batch.items():
        print(k, b.shape, sep='\t')
    print('\n')
```

### Expected behavior

For `num_workers` equal to 0 or 1 works fine as expected:

```
feature	torch.Size([16])
target	torch.Size([16])


feature	torch.Size([16])
target	torch.Size([16])
....
```

Expected to support workers >1.

### Environment info

Databricks 13.3 LTS ML runtime - Spark 3.4.1
pyspark==3.4.1
py4j==0.10.9.7
datasets==2.13.1 and also tested with datasets==2.14.6