LLM Assistant for Apache Spark

Installation

pip install spark-llm

Usage

Initialization

from spark_llm import SparkLLMAssistant

assistant = SparkLLMAssistant()
assistant.activate() # active partial functions for Spark DataFrame

Data Ingestion

auto_df = assistant.create_df("2022 USA national auto sales by brand")
auto_df.show(n=5)

Plot

auto_df.llm.plot()

To plot with an instruction:

auto_df.llm.plot("pie char for top 5 brands and the others' market shares")

DataFrame Transformation

auto_top_growth_df=auto_df.llm.transform("top brand with the highest growth")
auto_top_growth_df.show()

DataFrame Explanation

auto_top_growth_df.llm.explain()

In summary, this dataframe is retrieving the brand with the highest sales change in 2022 compared to 2021. It presents the results sorted by sales change in descending order and only returns the top result.

DataFrame Attribute Verification

auto_top_growth_df.llm.verify("expect sales change percentage to be between -100 to 100")

result: True

UDF Generation

@assistant.udf
def previous_years_sales(brand: str, current_year_sale: int, sales_change_percentage: float) -> int:
    """Calculate previous years sales from sales change percentage"""
    ...
    
spark.udf.register("previous_years_sales", previous_years_sales)
auto_df.createOrReplaceTempView("autoDF")

spark.sql("select brand as brand, previous_years_sales(brand, us_sales, sales_change_percentage) as 2021_sales from autoDF").show()

Cache

The SparkLLMAssistant supports a simple in-memory and persistent cache system. It keeps an in-memory staging cache that can be persisted through the commit() method. Cache lookup is always performed on the persistent cache only.

assistant.commit()

Refer to example.ipynb for more detailed usage examples.

Contributing

Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.

License

Licensed under the Apache License 2.0.