Databricks Connect + AutoDependencies + UDF + Streaming breaks

[i4ata]

I have this toy python script:

CATALOG = 'mpt_usd_gbl_bronze_dev' # <-- replace
SCHEMA = 'mpt_gbl_tco_fof' # <-- replace

src=f'{CATALOG}.{SCHEMA}.toy_source'
DEST = f'{CATALOG}.{SCHEMA}.toy_dest'

from databricks.connect import DatabricksSession, DatabricksEnv
from pyspark.sql import functions as F, SparkSession, DataFrame

env = DatabricksEnv().withAutoDependencies(upload_local=True, use_index=True)
spark: SparkSession = DatabricksSession.builder.serverless().withEnvironment(env).getOrCreate()

spark.sql(f'CREATE SCHEMA IF NOT EXISTS {CATALOG}.{SCHEMA}')
spark.sql(f'DROP TABLE IF EXISTS {SRC}')
spark.sql(f'DROP TABLE IF EXISTS {DEST}')

(
    spark.createDataFrame(
        [('Alice', 1), ('Bob', 2), ('Claire', 3),],
        schema='name string, value int',
    )
    .write.format('delta')
    .saveAsTable(SRC)
)


def test(batch_df: DataFrame, batch_id: int):
f = F.udf(lambda x: 'here')
(
    batch_df
    .withColumn('here', f(F.col('file_name')))
    .write.format('delta').mode('append')
    .saveAsTable(DEST)
)

(
    spark
    .readStream.format('delta').table(SRC)
    .writeStream.foreachBatch(test)
    .trigger(availableNow=True).start().awaitTermination()
)

It (1) initializes databricks-connect (2) loads dependencies (3) creates a UDF (4) runs a streaming query. I am getting this error:

Traceback (most recent call last):
File "/home/ivaylo-rusinov/tieto/magna/MPT_Cloud_Data_Analytics/MPT_UNS_Engineering_Pipeline_3S/PipelineConfigurations/MPT_GBL_Workspace/SourceAligned_GBL_TCO_FoF/tests3.py", line 40, in <module>
.trigger(availableNow=True).start().awaitTermination()
^^^^^^^
File "/home/ivaylo-rusinov/python_312_env/lib/python3.12/site-packages/pyspark/sql/connect/streaming/readwriter.py", line 793, in start
return self._start_internal(
^^^^^^^^^^^^^^^^^^^^^
File "/home/ivaylo-rusinov/python_312_env/lib/python3.12/site-packages/pyspark/sql/connect/streaming/readwriter.py", line 755, in _start_internal
_, properties, _ = self._session.client.execute_command(cmd)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ivaylo-rusinov/python_312_env/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 1880, in execute_command
data, _, metrics, observed_metrics, properties = self._execute_and_fetch(
^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ivaylo-rusinov/python_312_env/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 2512, in _execute_and_fetch
for response in self._execute_and_fetch_as_iterator(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ivaylo-rusinov/python_312_env/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 2266, in _execute_and_fetch_as_iterator
req = hook.on_execute_plan(req)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ivaylo-rusinov/python_312_env/lib/python3.12/site-packages/databricks/connect/auto_dependencies/hook.py", line 110, in on_execute_plan
obj = obj[0]
~~~^^^
TypeError: 'function' object is not subscriptable

Is this intended behavior? I assume this is quite a common workflow. How can I fix this?

[Jen_Evangelista]

Looking at the traceback, the failure happens inside the auto-dependencies hook (auto_dependencies/hook.py, on_execute_plan), not in your code. When the plan contains a foreachBatch function, the hook receives the bare function object and tries to index it (obj[0]), which raises "'function' object is not subscriptable". So it looks like a bug in how withAutoDependencies handles foreachBatch, not intended behavior.

Workarounds while it gets fixed:

1. Drop withAutoDependencies() for this job and declare dependencies explicitly with DatabricksEnv().withDependencies([...]). You lose the automatic detection, but the streaming query starts normally.

2. Keep auto-dependencies for interactive/batch work and run the streaming part as a job (wheel or notebook task), which is closer to how it will run in production anyway.

Two small things in the snippet, unrelated to the error: src is defined in lowercase but SRC is used later, and the UDF reads a file_name column that doesn't exist in toy_source, so the batch will fail once the hook issue is out of the way.

I'd also report it with your databricks-connect version (pip show databricks-connect), since this is a clear, reproducible case.

[Jen_Evangelista]

Quick follow-up: I reproduced the hook failure locally with the latest databricks-connect (19.2.0). foreachBatch sends its function as a PythonUDF message, but its command is just the pickled function, not the (func, return_type) tuple the hook expects. So the UDF inside your batch isn't the trigger: any foreachBatch combined with withAutoDependencies() hits this error, and it's still present in the latest version.

[i4ata]

Thanks for the reply and apologies for the typos in the snippet! I tried to get withDependencies() to work but as my UDF was requiring an external package, it was still failing. It can be replicated with the following snippet:

CATALOG = 'catalog' # <-- replace
SCHEMA = 'schema' # <-- replace

src=f'{CATALOG}.{SCHEMA}.toy_source'
DEST = f'{CATALOG}.{SCHEMA}.toy_dest'

from databricks.connect import DatabricksSession, DatabricksEnv
from pyspark.sql import functions as F, SparkSession, DataFrame
import openpyxl

env = DatabricksEnv().withDependencies('openpyxl')
spark: SparkSession = DatabricksSession.builder.serverless().withEnvironment(env).getOrCreate()

spark.sql(f'CREATE SCHEMA IF NOT EXISTS {CATALOG}.{SCHEMA}')
spark.sql(f'DROP TABLE IF EXISTS {SRC}')
spark.sql(f'DROP TABLE IF EXISTS {DEST}')

spark.createDataFrame([('Alice', 1), ('Bob', 2)], schema='name string, value int').write.format('delta').saveAsTable(SRC)
f = F.udf(lambda x: x + openpyxl.__version__)
def test(batch_df: DataFrame, batch_id: int):
  batch_df.withColumn('here', f(F.col('name'))).write.format('delta').mode('overwrite').saveAsTable(DEST)

spark.readStream.format('delta').table(SRC).writeStream.foreachBatch(test).trigger(availableNow=True).start().awaitTermination()

This results in:

File "/databricks/spark/python/_engine_pyspark.zip/_engine_pyspark/sql/connect/streaming/worker/foreach_batch_worker.py", line 173, in main
process(df_ref_id, int(batch_id))
File "/databricks/spark/python/_engine_pyspark.zip/_engine_pyspark/sql/connect/streaming/worker/foreach_batch_worker.py", line 77, in process
func(batch_df, batch_id)
File "/home/ivaylo-rusinov/databricks_forums/main.py", line 21, in test
File "/databricks/python3/lib/python3.12/site-packages/pyspark/sql/connect/readwriter.py", line 828, in saveAsTable
_, _, ei = self._spark.client.execute_command(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/databricks/python3/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 1812, in execute_command
data, _, metrics, observed_metrics, properties = self._execute_and_fetch(
^^^^^^^^^^^^^^^^^^^^^^^^
File "/databricks/python3/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 2411, in _execute_and_fetch
for response in self._execute_and_fetch_as_iterator(
File "/databricks/python3/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 2387, in _execute_and_fetch_as_iterator
self._handle_error(error, req.operation_id)
File "/databricks/python3/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 2785, in _handle_error
self._handle_rpc_error(error)
File "/databricks/python3/lib/python3.12/site-packages/pyspark/sql/connect/client/core.py", line 2882, in _handle_rpc_error
raise convert_exception(
File "/databricks/python3/lib/python3.12/site-packages/pyspark/cloudpickle/cloudpickle.py", line 475, in subimport
__import__(name)
^^^^^^^^^^^^^^^
pyspark.errors.exceptions.connect.PythonException: 
An exception was thrown from the Python worker. Please see the stack trace below.
[PYTHON_EXCEPTION] An exception was thrown from the Python worker: Traceback (most recent call last):
File "/databricks/python3/lib/python3.12/site-packages/pyspark/cloudpickle/cloudpickle.py", line 475, in subimport
__import__(name)
ModuleNotFoundError: No module named 'openpyxl'

I am also using databricks-connect 19.2.0. This runs fine outside of `foreachBatch`. I guess the functionality is just too limited.