from pyspark.sql import SparkSession
from sklearn import datasets
from sklearn.neighbors import KNeighborsClassifier
import mlflow
from mlflow.models import infer_signature

spark = SparkSession.builder.getOrCreate()

#データセットをロードする
#方法1：機械学習ライブラリを使用してデータセットをロードする
X, y = datasets.load_iris(as_frame=True, return_X_y=True)
#方法2：tencentcloud-dlc-connectorを使用してDLCのデータをロードする
#ドライバをインストールする
!pip install tencentcloud-dlc-connector
!pip  install --upgrade 'sqlalchemy<2.0'
#バージョンをインストールする
!pip install --upgrade pandas==2.2.3
!pip install numpy
!pip install matplotlib
import pandas as pd
import numpy as np
import tdlc_connector
from tdlc_connector import constants
mlflow.sklearn.autolog()
# tdlc-connectorを使用してテーブル方式でアクセスします
conn = tdlc_connector.connect(region="ap-***", #正しいアドレスを入力してください、例：ap-Singapore,ap-Shanghai
    secret_id="*******", 
    secret_key="*******",
    engine="your engine",#購入したエンジン名を入力してください
    resource_group=None,
    engine_type=constants.EngineType.AUTO, 
    result_style=constants.ResultStyles.LIST,
    download=True
    )
query = """
SELECT `sepal.length`, `sepal.width`,`petal.length`,`petal.width`,species FROM at_database_testnotebook.demo_test_sklearn
"""
iris = pd.read_sql(query, conn)
spark_iris = spark.createDataFrame(iris)
#特徴列とターゲット列の分割
feature_cols = ["sepal_length", "sepal_width", "petal_length", "petal_width"]X = spark_iris.select(feature_cols)
X = spark_iris.select(feature_cols)
y = spark_iris.select("species")

#K近傍法を使用した分類
model = KNeighborsClassifier()
model.fit(X, y)
predictions = model.predict(X)
signature = infer_signature(X, predictions)

with mlflow.start_run():
    model_info = mlflow.sklearn.log_model(model, artifact_path="model", signature=signature)
    
infer_spark_df = spark.createDataFrame(X)
pyfunc_udf = mlflow.pyfunc.spark_udf(spark, model_info.model_uri)
result = infer_spark_df.select(pyfunc_udf(*X.columns).alias("predictions")).toPandas()
print(result)