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.