"""Costes de seguros médicos con PySpark: EDA con pandas-on-Spark y regresión lineal con Spark ML. Datos: insurance.csv (1.338 beneficiarios), de stedy/Machine-Learning-with-R-datasets. Reproduce el notebook del proyecto (mismo split, mismos pipelines y misma búsqueda de hiperparámetros) y genera data/seguros.json para la web. Uso (Spark 4 necesita Java 17 o 21; con Java 23 añade -Djava.security.manager=allow): JAVA_TOOL_OPTIONS="-Djava.security.manager=allow" python analysis/insurance_pyspark.py raw/insurance.csv """ import json import sys from pathlib import Path import numpy as np import pyspark.pandas as ps from hyperopt import STATUS_OK, Trials, fmin, hp, tpe from pyspark.ml import Pipeline from pyspark.ml.evaluation import RegressionEvaluator from pyspark.ml.feature import OneHotEncoder, StandardScaler, StringIndexer, VectorAssembler from pyspark.ml.regression import LinearRegression from pyspark.ml.tuning import CrossValidator, ParamGridBuilder from pyspark.sql import SparkSession from pyspark.sql import functions as F ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "data" / "seguros.json" def main(path): spark = SparkSession.builder.appName("seguros").getOrCreate() spark.sparkContext.setLogLevel("ERROR") spark.conf.set("spark.sql.ansi.enabled", "false") # pandas-on-Spark no funciona con ANSI activado # ---------- EDA con pandas-on-Spark ---------- df = ps.read_csv(str(path)) df["is_smoker"] = (df["smoker"] == "yes").astype(int) df["is_obese"] = (df["bmi"] >= 30).astype(int) df["smoker_obese_interaction"] = df["is_smoker"] * df["is_obese"] corr = df[["age", "bmi", "children", "is_smoker", "is_obese", "smoker_obese_interaction", "charges"]] \ .to_pandas().corr()["charges"].drop("charges").round(3) g = df.groupby(["smoker", "is_obese"])["charges"] # SeriesGroupBy.agg no existe en pandas-on-Spark grupos = g.count().to_pandas().to_frame("count").join(g.mean().to_pandas().rename("mean")) \ .join(g.median().to_pandas().rename("median")) q1, q3 = df["charges"].quantile(0.25), df["charges"].quantile(0.75) outliers = int(((df["charges"] < q1 - 1.5 * (q3 - q1)) | (df["charges"] > q3 + 1.5 * (q3 - q1))).sum()) puntos = df[["age", "charges", "is_smoker", "is_obese"]].to_pandas() # Mismas columnas auxiliares que el notebook: cambian el orden interno de las filas y, # con él, el reparto de los folds de la validación cruzada df["perfil_riesgo"] = "No Fumador" df["perfil_riesgo"] = df["perfil_riesgo"].mask((df["smoker"] == "yes") & (df["bmi"] < 30), "Fumador + Peso Normal") df["perfil_riesgo"] = df["perfil_riesgo"].mask((df["smoker"] == "yes") & (df["bmi"] >= 30), "Fumador + Obeso (BMI >= 30)") df["is_obese"] = df["bmi"] >= 30 df["obeso_texto"] = "Peso Normal" df["obeso_texto"] = df["obeso_texto"].mask(df["bmi"] >= 30, "Obeso (BMI>=30)") # ---------- Modelado con Spark ML ---------- data = df.to_spark().withColumn("log_charges", F.log1p("charges")) # Split ANTES de ajustar ningún transformador, para que el test no influya en el preprocesado train, test = data.randomSplit([0.8, 0.2], seed=42) train, test = train.repartition(8), test.repartition(8) indexer = StringIndexer(inputCols=["sex", "smoker", "region"], outputCols=["sex_indexed", "smoker_indexed", "region_indexed"]) ohe = OneHotEncoder(inputCol="region_indexed", outputCol="region_ohe", dropLast=True) num = VectorAssembler(inputCols=["age", "bmi", "children"], outputCol="numeric_features") scaler = StandardScaler(inputCol="numeric_features", outputCol="numeric_features_scaled", withStd=True, withMean=True) base_cols = ["sex_indexed", "smoker_indexed", "region_ohe", "numeric_features_scaled"] final = VectorAssembler(inputCols=base_cols, outputCol="features") final_inter = VectorAssembler(inputCols=base_cols + ["smoker_obese_interaction"], outputCol="features") stages = [indexer, ohe, num, scaler] resultados = [] def metricas(model, nombre, log_target): fila = {"modelo": nombre} for conjunto, frame in (("train", train), ("test", test)): pred = model.transform(frame).withColumn( "pred_usd", F.expm1("prediction") if log_target else F.col("prediction")) for m in ("r2", "rmse", "mae"): ev = RegressionEvaluator(labelCol="charges", predictionCol="pred_usd", metricName=m) fila[f"{m}_{conjunto}"] = round(ev.evaluate(pred), 4 if m == "r2" else 2) resultados.append(fila) print(fila) # 1 y 2: regresión lineal sobre charges y sobre log(charges) m1 = Pipeline(stages=stages + [final, LinearRegression(labelCol="charges", regParam=0.01)]).fit(train) metricas(m1, "1", log_target=False) m2 = Pipeline(stages=stages + [final, LinearRegression(labelCol="log_charges", regParam=0.01)]).fit(train) metricas(m2, "2", log_target=True) # Hyperopt con una validación DENTRO de train (el test no se toca) sub_tr, sub_val = train.randomSplit([0.8, 0.2], seed=42) prep = Pipeline(stages=stages + [final]).fit(sub_tr) sub_tr_p, sub_val_p = prep.transform(sub_tr).persist(), prep.transform(sub_val).persist() ev_log = RegressionEvaluator(labelCol="log_charges", predictionCol="prediction", metricName="rmse") def objective(params): model = LinearRegression(labelCol="log_charges", regParam=params["regParam"], elasticNetParam=params["elasticNetParam"]).fit(sub_tr_p) return {"loss": ev_log.evaluate(model.transform(sub_val_p)), "status": STATUS_OK} best = fmin(objective, {"regParam": hp.loguniform("regParam", -9.0, 3.5), "elasticNetParam": hp.uniform("elasticNetParam", 0.0, 0.5)}, algo=tpe.suggest, max_evals=20, trials=Trials(), rstate=np.random.default_rng(42), show_progressbar=False) reg_values = sorted({0.001, 0.01, 0.1, 1.0, round(float(best["regParam"]), 4)}) en_values = sorted({0.0, 0.25, 0.5, round(float(best["elasticNetParam"]), 4)}) def cross_validar(assembler, label): lr = LinearRegression(featuresCol="features", labelCol=label) grid = ParamGridBuilder().addGrid(lr.regParam, reg_values).addGrid(lr.elasticNetParam, en_values).build() cv = CrossValidator(estimator=Pipeline(stages=stages + [assembler, lr]), estimatorParamMaps=grid, evaluator=RegressionEvaluator(labelCol=label, metricName="rmse"), numFolds=3, seed=42) return cv.fit(train).bestModel # 3: log + CV · 4: log + CV + interacción · 5: dólares + CV + interacción metricas(cross_validar(final, "log_charges"), "3", log_target=True) metricas(cross_validar(final_inter, "log_charges"), "4", log_target=True) m5 = cross_validar(final_inter, "charges") metricas(m5, "5", log_target=False) lr5 = m5.stages[-1] # Real frente a predicho del mejor modelo en test pred5 = m5.transform(test).select("charges", "prediction", "is_smoker", "is_obese").toPandas() out = { "kpis": {"filas": int(len(puntos)), "train": int(train.count()), "test": int(test.count()), "outliers": outliers, "hyperopt": {k: round(float(v), 4) for k, v in best.items()}, "best_reg": float(lr5.getOrDefault("regParam")), "best_en": float(lr5.getOrDefault("elasticNetParam"))}, "corr": corr.to_dict(), "grupos": [{"smoker": s, "obese": int(o), "n": int(r["count"]), "media": round(r["mean"]), "mediana": round(r["median"])} for (s, o), r in grupos.iterrows()], "puntos": puntos.round(0).astype(int).values.tolist(), "modelos": resultados, "pred": pred5.round(0).astype(int).values.tolist(), } OUT.parent.mkdir(exist_ok=True) OUT.write_text(json.dumps(out, ensure_ascii=False, separators=(",", ":"))) print(json.dumps(out["kpis"], indent=2), out["corr"], out["grupos"], sep="\n") spark.stop() if __name__ == "__main__": main(sys.argv[1] if len(sys.argv) > 1 else ROOT / "raw" / "insurance.csv")