Quick tip: Using Apache Spark Structured Streaming with SingleStore Notebooks

!pip cache purge --quiet!conda install -y --quiet -c conda-forge openjdk pyspark

!pip install openai==0.28 --quiet!pip install nltk --quiet

import getpassimport openaios.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

import osos.makedirs("jars", exist_ok = True)os.makedirs("data", exist_ok = True)

import requestsdef download_jar(url, destination):    response = requests.get(url)    with open(destination, "wb") as f:        f.write(response.content)jar_urls = [    ("https://repo1.maven.org/maven2/com/singlestore/singlestore-jdbc-client/1.2.1/singlestore-jdbc-client-1.2.1.jar", "jars/singlestore-jdbc-client-1.2.1.jar"),    ("https://repo1.maven.org/maven2/com/singlestore/singlestore-spark-connector_2.12/4.1.5-spark-3.5.0/singlestore-spark-connector_2.12-4.1.5-spark-3.5.0.jar", "jars/singlestore-spark-connector_2.12-4.1.5-spark-3.5.0.jar"),    ("https://repo1.maven.org/maven2/org/apache/commons/commons-dbcp2/2.12.0/commons-dbcp2-2.12.0.jar", "jars/commons-dbcp2-2.12.0.jar"),    ("https://repo1.maven.org/maven2/org/apache/commons/commons-pool2/2.12.0/commons-pool2-2.12.0.jar", "jars/commons-pool2-2.12.0.jar"),    ("https://repo1.maven.org/maven2/io/spray/spray-json_3/1.3.6/spray-json_3-1.3.6.jar", "jars/spray-json_3-1.3.6.jar")]for url, destination in jar_urls:    download_jar(url, destination)print("JAR files downloaded successfully")

from pyspark.sql import SparkSession# Create a Spark sessionspark = (SparkSession             .builder             .config("spark.jars", ",".join([destination for _, destination in jar_urls]))             .appName("Spark Streaming Test")             .getOrCreate()        )spark.sparkContext.setLogLevel("ERROR")

import nltkimport randomfrom nltk.corpus import wordnet as wnfrom nltk.tokenize import word_tokenize# Download NLTKnltk.download("punkt")nltk.download("averaged_perceptron_tagger")nltk.download("wordnet")nltk.download("omw")# Define the directory to save the filesoutput_dir = "data"# Generate meaningful sentencesdef generate_meaningful_sentence():    # Choose a random set of synonyms from WordNet    synset = random.choice(list(wn.all_synsets()))    # Generate a sentence    definition = synset.definition()    tokens = word_tokenize(definition)    # Capitalise the first word and end with a period    tokens[0] = tokens[0].capitalize()    tokens[-1] = tokens[-1] + "."    return " ".join(tokens)# Number of files to generatenum_files = 5# Number of sentences in each filenum_sentences_per_file = 1# Generate text filesfor i in range(num_files):    file_path = os.path.join(output_dir, f"file_{i+1}.txt")    with open(file_path, "w") as file:        for _ in range(num_sentences_per_file):            # Generate a meaningful sentence            sentence = generate_meaningful_sentence()            file.write(sentence + "
")

%%bashfor file in data/*.txt; do    echo "File: $file"    cat "$file"    echo "----------------------"done

File: data/file_1.txtResembling steel in hardness.----------------------File: data/file_2.txtAn instruction that is either not recognized by an operating system or is in violation of the procedural rules.----------------------File: data/file_3.txtDestructive to the spirit or soul.----------------------File: data/file_4.txtThornbills.----------------------File: data/file_5.txtA small mass of soft material.----------------------

host = "<HOST>"password = "<PASSWORD>"port = "3306"cluster = host + ":" + port

spark.conf.set("spark.datasource.singlestore.ddlEndpoint", cluster)spark.conf.set("spark.datasource.singlestore.user", "admin")spark.conf.set("spark.datasource.singlestore.password", password)spark.conf.set("spark.datasource.singlestore.disablePushdown", "false")

DROP DATABASE IF EXISTS spark_demo;CREATE DATABASE IF NOT EXISTS spark_demo;

USE spark_demo;DROP TABLE IF EXISTS streaming;CREATE TABLE IF NOT EXISTS streaming (     value TEXT,     file_name TEXT,     embedding VECTOR(1536) NOT NULL);

from pyspark.sql.functions import input_file_name, udffrom pyspark.sql.types import StringType# Generate embeddings for textdef generate_embeddings(text):    openai.api_key = os.environ.get("OPENAI_API_KEY")    # Generate embeddings for text using OpenAI    return openai.Embedding.create(        input = text,        engine = "text-embedding-3-small"    ).data[0].embedding# Register the function as a UDFgenerate_embeddings_udf = udf(generate_embeddings, StringType())

import timeinput_dir = output_dir# Read from the directorydf = (spark.readStream    .format("text")    .option("path", input_dir)    .load()    .withColumn("file_name", input_file_name()))# Apply the function to the DataFrame to generate embeddings for each rowdf_with_embeddings = df.withColumn("embedding", generate_embeddings_udf("value"))# Write each batch of data to SingleStoredef write_to_singlestore(df_with_embeddings, epoch_id):    (df_with_embeddings.write         .format("singlestore")         .option("loadDataCompression", "LZ4")         .mode("append")         .save("spark_demo.streaming")    )# Write the streaming DataFrame to SingleStore using foreachBatchquery = (df_with_embeddings.writeStream    .foreachBatch(write_to_singlestore)    .start())# Wait for the query to finish processingwhile query.isActive:    time.sleep(1)    if not query.status["isDataAvailable"]:        query.stop()

USE spark_demo;SELECT    SUBSTR(value, 1, 30) AS value,    SUBSTR(file_name, LENGTH(file_name) - 9) AS file_name,    SUBSTR(embedding, 1, 50) AS embeddingFROM streaming;

spark.stop()