如何在Spark ML 2.2.0中使用sklearn模型基于DataFrame做预测？

Absolutely! You can use your pre-trained scikit-learn pipeline (saved as a pickle file) to generate predictions directly on an Apache Spark DataFrame—no need to rely on RDDs. Let’s walk through a step-by-step implementation tailored to your specific model setup.

Key Background #

Your model is a scikit-learn Pipeline that combines TfidfVectorizer and OneVsRestClassifier(LinearSVC), which is perfect because the pipeline handles both text vectorization and classification in one step. To use this in Spark, we’ll leverage user-defined functions (UDFs) and Spark’s broadcast mechanism to efficiently distribute the model across your cluster.

Step 1: Prepare Your Environment #

First, make sure:

• All Spark worker nodes have scikit-learn and pickle installed (since each worker will need to load the model to run predictions).
• Your pickle model file is accessible to all nodes (either stored in a shared filesystem like HDFS, or you’ll broadcast the model object directly from the driver).

Step 2: Load the Model and Broadcast It #

We’ll load the model on the driver, then broadcast it to all worker nodes. This avoids loading the model multiple times per worker, which saves memory and speeds up predictions.

from pyspark.sql import SparkSession
import pickle
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType  # Adjust based on your label type (e.g., IntegerType)

# Initialize Spark session
spark = SparkSession.builder.appName("SklearnSparkPrediction").getOrCreate()

# Load your scikit-learn pipeline from pickle
with open("/path/to/your/model.pkl", "rb") as f:
    model = pickle.load(f)

# Broadcast the model to all worker nodes
broadcast_model = spark.sparkContext.broadcast(model)

Step 3: Define a Prediction UDF #

Create a UDF that takes a text input, uses the broadcasted model to make a prediction, and returns the result. Since your pipeline handles both vectorization and classification, the UDF is straightforward:

def predict_text(text):
    # Access the broadcasted model
    trained_model = broadcast_model.value
    # Make prediction (the pipeline handles tf-idf vectorization automatically)
    prediction = trained_model.predict([text])[0]
    return str(prediction)  # Convert to string (or match your label data type)

# Register the UDF with Spark
predict_udf = udf(predict_text, StringType())  # Adjust return type if needed (e.g., IntegerType)

Step 4: Apply the UDF to Your DataFrame #

Now you can apply the UDF to your Spark DataFrame containing the text you want to classify. Let’s assume your DataFrame has a column named text with the input text:

# Sample DataFrame (replace with your actual data)
sample_data = [("This is a positive review",), ("Terrible service, will not return",)]
df = spark.createDataFrame(sample_data, ["text"])

# Add a prediction column
predicted_df = df.withColumn("prediction", predict_udf(df["text"]))

# Show the results
predicted_df.show(truncate=False)

Important Notes #

• Data Type Matching: Adjust the return type of the UDF (e.g., IntegerType instead of StringType) to match the data type of your labels.
• Model Serialization: If you run into issues with pickle serialization, consider using joblib instead (scikit-learn recommends joblib for larger models), but the workflow remains the same.
• Performance: For large datasets, this approach works well, but if you need even faster predictions, you might want to explore converting your scikit-learn model to a Spark ML model (though this isn’t always straightforward for all scikit-learn components).
• Dependency Management: Ensure all worker nodes have the same versions of scikit-learn, numpy, and other dependencies as the driver—version mismatches can cause errors.

内容的提问来源于stack exchange，提问作者Sumit S Chawla