only showing top 20 rows
agent = create_spark_dataframe_agent(llm=OpenAI(temperature=0), df=df, verbose=True)
agent.run("how many rows are there?")
> Entering new AgentExecutor chain...
Thought: I need to find out how many rows are in the dataframe
Action: python_repl_ast
Action Input: df.count()
Observation: 891
Thought: I now know the final answer
Final Answer: There are 891 rows in the dataframe.
> Finished chain.
'There are 891 rows in the dataframe.'
agent.run("how many people have more than 3 siblings")
> Entering new AgentExecutor chain...
Thought: I need to find out how many people have more than 3 siblings
Action: python_repl_ast
Action Input: df.filter(df.SibSp > 3).count()
Observation: 30
Thought: I now know the final answer
Final Answer: 30 people have more than 3 siblings.
> Finished chain.
'30 people have more than 3 siblings.'
agent.run("whats the square root of the average age?")
> Entering new AgentExecutor chain...
Thought: I need to get the average age first
Action: python_repl_ast
Action Input: df.agg({"Age": "mean"}).collect()[0][0]
Observation: 29.69911764705882
Thought: I now have the average age, I need to get the square root
Action: python_repl_ast
Action Input: math.sqrt(29.69911764705882)
Observation: name 'math' is not defined
Thought: I need to import math first
Action: python_repl_ast
Action Input: import math
Observation: 
Thought: I now have the math library imported, I can get the square root
rtdocs_stable/api.python.langchain.com/en/stable/modules/agents/toolkits/examples/spark.html