[sklearn in spark] spark분산환경을 이용한 모델 예측

sklearn 모델을 spark 환경에서 분산 인퍼런스하기!!!

- pandas udf를 활용

: spark 드라이버가 각 worker에게 모델을 전달하고, 각 worker는 해당 모델을 이용해 pandas udf에서 spark dataframe을 pandas dataframe으로 변환해 예측하고 이를 다시 spark dataframe형태로 리턴

 

코드예시

from pyspark.sql.functions import col
from sklearn.linear_model import LinearRegression

# 데이터 생성
data = [("row1", 1.0, 2.0, 3.0),
        ("row2", 4.0, 5.0, 6.0)]

columns = ["id", "feature1", "feature2", "feature3"]
df = spark.createDataFrame(data, columns)

# Scikit-Learn 모델 정의 및 학습
features = ["feature2", "feature3"]
X_train = np.array(df.select(*features).collect())
y_train = np.array(df.select("feature1").collect())
model = LinearRegression()
model.fit(X_train, y_train)

# Python 함수로 모델 추론 정의
def predict_udf(*features):
    features_array = np.array(features).reshape(1, -1)
    prediction = model.predict(features_array)
    return float(prediction[0])

# UDF로 Python 함수 등록
predict_udf = spark.udf.register("predict_udf", predict_udf)

# UDF를 사용하여 예측 컬럼 생성 (별표(*)를 사용하여 모든 feature를 선택)
result_df = df.withColumn("prediction", predict_udf(*features))

# 결과 확인
result_df.select("id", "prediction").show(truncate=False)