ck7007
Contributor II

@Yousry_Ibrahim This is a known limitation with PySpark custom data sources in shared access mode. The issue is that custom data sources serialize differently from UDFs.

Root Cause

Custom data sources use cloudpickle serialization, which doesn't properly capture the sys.path modifications in shared clusters. UDFs work because they use a different serialization path.

Solution 1: Add Module to Spark Files

# Add the module file explicitly to the Spark context
spark.sparkContext.addPyFile("/Workspace/path/to/module2.py")

# Now import and register
from module2 import DummyDataSource
spark.dataSource.register(DummyDataSource)
spark.read.format("dummy"). load(). display()

Solution 2: Package as an Init-Script

Create an init script that adds your modules to the Python path on all nodes:
#!/bin/bash
# cluster-init.sh
echo "export PYTHONPATH=/Workspace/your_modules:$PYTHONPATH" >> /databricks/spark/conf/spark-env.sh

Solution 3: Inline the Data Source (Temporary Fix)

For testing, define the data source directly in the notebook:
# Define the classes in the notebook itself
exec(open('/Workspace/path/to/module2.py').read())
spark.dataSource.register(DummyDataSource)

Why This Happens

  • UDFs: Executed in Python worker processes that inherit sys.path
  • Data Sources: Serialized at the driver and deserialized at executors without sys.path context
  • Shared Mode: Additional isolation prevents path propagation

The addPyFile approach is the most reliable for shared clusters. It ensures the module is distributed to all executors before deserialization.

Have you tried Solution 1? It should work immediately without a cluster restart.

View solution in original post