<jupyter_start><jupyter_text>This is a tutorial of how to bring Pandas and Python code to Spark. We'll compare the traditional way with Fugue.<jupyter_code>import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression

X = pd.DataFrame({"x_1": [1, 1, 2, 2], "x_2":[1, 2, 2, 3]})
y = np.dot(X, np.array([1, 2])) + 3
reg = LinearRegression().fit(X, y)<jupyter_output><empty_output><jupyter_text>Create the prediction function and test it<jupyter_code>def predict(df: pd.DataFrame, model: LinearRegression) -> pd.DataFrame:
    return df.assign(predicted=model.predict(df))

input_df = pd.DataFrame({"x_1": [3, 4, 6, 6], "x_2":[3, 3, 6, 6]})
predict(input_df.copy(), reg)
from fugue import transform
from fugue_spark import SparkExecutionEngine

result = transform(
    input_df,
    predict,
    schema="*,predicted:double",
    params={"model": reg},
    engine=SparkExecutionEngine()
)
result.show()
from typing import Iterator, Any, Union
from pyspark.sql.types import StructType, StructField, DoubleType
from pyspark.sql import DataFrame, SparkSession

spark_session = SparkSession.builder.getOrCreate()

def predict_wrapper(dfs: Iterator[pd.DataFrame], model):
    for df in dfs:
        yield predict(df, model)

def run_predict(input_df: Union[DataFrame, pd.DataFrame], model):
    # conversion
    if isinstance(input_df, pd.DataFrame):
        sdf = spark_session.createDataFrame(input_df.copy())
    else:
        sdf = input_df.copy()

    schema = StructType(list(sdf.schema.fields))
    schema.add(StructField("predicted", DoubleType()))
    return sdf.mapInPandas(lambda dfs: predict_wrapper(dfs, model),
                           schema=schema)

result = run_predict(input_df.copy(), reg)
result.show()<jupyter_output>+---+---+---------+
|x_1|x_2|predicted|
+---+---+---------+
|  3|  3|     12.0|
|  4|  3|     13.0|
|  6|  6|     21.0|
|  6|  6|     21.0|
+---+---+---------+
{'fraction_non_alphanumeric': 0.15745568300312826, 'fraction_numerical': 0.02346193952033368, 'mean_word_length': 3.8095238095238093, 'pattern_counts': {'":': 5, '<': 8, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '335685', 'license': 'MIT', 'n_tokens_mistral': 711, 'n_tokens_neox': 675, 'n_words': 163, 'path': 'blogs/fugue-transform/Untitled.ipynb', 'repo_name': 'kvnkho/demos', 'size': '1918'}