Databricks上使用LangChain遇Spark未安装ImportError（已装PySpark）

问题描述 #

在Databricks环境执行以下LangChain代码时：

from langchain_experimental.agents.agent_toolkits.spark.base import create_spark_dataframe_agent
agent = create_spark_dataframe_agent(llm=OpenAI(temperature=0), df=graphframe, verbose=True)
agent.run('How many rows have premium cuts')

触发ImportError: Spark is not installed. run pip install pyspark错误，完整报错堆栈如下：

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
File <command-2122371644497407>, line 3
      1 from langchain_experimental.agents.agent_toolkits.spark.base import create_spark_dataframe_agent
----> 3 agent = create_spark_dataframe_agent(llm=OpenAI(temperature=0), df=graphframe, verbose=True)
      4 agent.run('How many rows have premium cuts')

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-64863c5a-c177-4f58-adac-d5e56cd70936/lib/python3.10/site-packages/langchain_experimental/agents/agent_toolkits/spark/base.py:50, in create_spark_dataframe_agent(llm, df, callback_manager, prefix, suffix, input_variables, verbose, return_intermediate_steps, max_iterations, max_execution_time, early_stopping_method, agent_executor_kwargs, **kwargs)
     47 """Construct a Spark agent from an LLM and dataframe."""
     49 if not _validate_spark_df(df) and not _validate_spark_connect_df(df):
---&gt; 50     raise ImportError("Spark is not installed. run `pip install pyspark`.")
     52 if input_variables is None:
     53     input_variables = ["df", "input", "agent_scratchpad"]

ImportError: Spark is not installed. run `pip install pyspark`.

用户确认集群库已完整安装PySpark，仍出现该错误，以下是针对性解决办法：

解决步骤 #

1. 检查传入的DataFrame类型 #

报错根源是LangChain的_validate_spark_df和_validate_spark_connect_df函数未识别出graphframe为Spark DataFrame：

• 执行print(type(graphframe))验证类型，正确输出应为<class 'pyspark.sql.dataframe.DataFrame'>；
• 若graphframe是GraphFrame类型，需提取其内部的Vertices或Edges DataFrame（如graphframe.vertices、graphframe.edges）传入create_spark_dataframe_agent函数。

2. 对齐Python环境 #

Databricks存在集群全局环境与Notebook专属环境的差异，LangChain可能运行在未关联PySpark的环境中：

• 在Notebook开头执行%pip install pyspark，确保当前Notebook环境安装PySpark；
• 检查集群Python版本与Notebook环境版本是否一致，避免版本不兼容导致的导入失效。

3. 手动验证Spark导入有效性 #

在执行LangChain代码前，先测试PySpark是否能正常导入：

import pyspark
from pyspark.sql import DataFrame
print("PySpark导入成功")

若此步骤报错，需检查集群Spark版本配置，或重启集群确保库加载生效。

4. 强制转换为Spark DataFrame（临时方案） #

若上述方法无效，可手动将传入对象转换为标准Spark DataFrame：

from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()

# 若graphframe是Pandas DataFrame
spark_df = spark.createDataFrame(graphframe)
# 若graphframe是GraphFrame，提取对应DataFrame
# spark_df = graphframe.vertices

agent = create_spark_dataframe_agent(llm=OpenAI(temperature=0), df=spark_df, verbose=True)
agent.run('How many rows have premium cuts')

内容的提问来源于stack exchange，提问作者Isa Khan