from databricks.feature_engineering import FeatureEngineeringClient, FeatureLookup
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import f1_score
import mlflow
fe = FeatureEngineeringClient()
# ── Step 1: 피처 계산 및 테이블 생성 ──
customer_features_df = (
spark.table("catalog.schema.transactions")
.groupBy("customer_id")
.agg(
F.count("*").alias("total_transactions"),
F.avg("amount").alias("avg_amount"),
F.stddev("amount").alias("stddev_amount"),
F.countDistinct("merchant_id").alias("unique_merchants")
)
)
fe.create_table(
name="catalog.ml_features.customer_features",
primary_keys=["customer_id"],
df=customer_features_df,
description="고객별 거래 통계 피처"
)
# ── Step 2: 학습 데이터 생성 (FeatureLookup) ──
labels_df = spark.table("catalog.schema.fraud_labels")
training_set = fe.create_training_set(
df=labels_df,
feature_lookups=[
FeatureLookup(
table_name="catalog.ml_features.customer_features",
lookup_key="customer_id"
)
],
label="is_fraud"
)
training_df = training_set.load_df().toPandas()
# ── Step 3: 모델 학습 ──
X = training_df.drop(columns=["is_fraud", "customer_id"])
y = training_df["is_fraud"]
model = GradientBoostingClassifier(n_estimators=200, learning_rate=0.1)
model.fit(X, y)
# ── Step 4: 모델 저장 (피처 메타데이터 포함) ──
with mlflow.start_run(run_name="fraud-with-feature-store"):
fe.log_model(
model=model,
artifact_path="model",
flavor=mlflow.sklearn,
training_set=training_set, # 피처 조회 정보가 모델에 기록됨
registered_model_name="catalog.schema.fraud_model_with_features"
)
# ── Step 5: 서빙 시 자동 피처 조회 ──
# 엔드포인트에 배포 후, 기본 키만 전달하면 피처는 자동 조회됩니다
# 요청 예시: {"dataframe_records": [{"customer_id": "C12345"}]}
# → customer_features 테이블에서 피처를 자동으로 가져와 예측 수행