ML su Snowflake su larga scala con Snowpark Python e XGBoost

Nell'ultimo anno e mezzo, il nostro team ha lavorato con centinaia di clienti costruendo su Snowflake con Snowpark e Python (a proposito, dai un'occhiata a questo post del mio collega Caleb Baechtold che raccoglie tutti gli insegnamenti sull'operatività di Snowpark Python in produzione) e un la domanda che ricevo spesso è: "ma abbiamo davvero big data, come funziona su larga scala?".

Molti degli esempi/demo che troverai online sono ottimi per mostrare semplici esempi inventati che mostrano le funzionalità di base, ma pochi o nessuno mostrano esempi di come funziona nel contesto di problemi su larga scala aziendale.

TPC-DS , è un utile set di dati su larga scala che generalmente rappresenta uno schema di dati aziendali generici (non sono qui per sostenere che TPC-DS o qualsiasi set di dati sia la fine di tutti i benchmark di qualsiasi piattaforma/o strumento, ho solo penso che possa essere un set di dati utile come punto di partenza per testare le cose su larga scala). Snowflake rende disponibile TPC-DS in ogni account Snowflake come condivisione dati, sia in un'edizione da 10 TB che in una da 100 TB (a proposito, poiché è esposto come condivisione dati tu come utente non paghi alcun costo per l'archiviazione . Solo per qualsiasi potenziale calcolo utilizzato per interrogarlo). L'edizione da 100 TB ha oltre 560 miliardi di righe nelle tabelle dei fatti. L'edizione da 10 TB, 56 miliardi.

Possiamo utilizzare Snowpark Python per creare una soluzione ML relativamente semplice per un problema aziendale comune di un'attività ipotetica come TPC che è "Voglio essere in grado di prevedere il valore della vita dei miei clienti su tutti i canali di vendita". Utilizzeremo l' API Snowpark Python Dataframe per eseguire la preparazione dei dati/l'ingegnerizzazione delle funzionalità, le stored procedure e i magazzini ottimizzati di Snowpark per l'addestramento e le UDF in batch per l'inferenza, il tutto senza che i dati debbano mai lasciare Snowflake utilizzando le risorse di elaborazione e la capacità di ridimensionamento di Snowflake.

Iniziamo con il nostro codice di preparazione dei dati/ingegneria delle funzionalità. Abbastanza semplicemente, aggregheremo le vendite per cliente su tutti i canali . Lo uniremo quindi alle nostre tabelle delle dimensioni dei clienti per ottenere potenziali caratteristiche di interesse.

store_sales_agged = store_sales.group_by('ss_customer_sk').agg(F.sum('ss_sales_price').as_('total_sales'))
web_sales_agged = web_sales.group_by('ws_bill_customer_sk').agg(F.sum('ws_sales_price').as_('total_sales'))
catalog_sales_agged = catalog_sales.group_by('cs_bill_customer_sk').agg(F.sum('cs_sales_price').as_('total_sales'))
store_sales_agged = store_sales_agged.rename('ss_customer_sk', 'customer_sk')
web_sales_agged = web_sales_agged.rename('ws_bill_customer_sk', 'customer_sk')
catalog_sales_agged = catalog_sales_agged.rename('cs_bill_customer_sk', 'customer_sk')
total_sales = store_sales_agged.union_all(web_sales_agged)
total_sales = total_sales.union_all(catalog_sales_agged)
total_sales = total_sales.group_by('customer_sk').agg(F.sum('total_sales').as_('total_sales'))
customer = customer.select('c_customer_sk','c_current_hdemo_sk', 'c_current_addr_sk', 'c_customer_id', 'c_birth_year')
customer = customer.join(address.select('ca_address_sk', 'ca_zip'), customer['c_current_addr_sk'] == address['ca_address_sk'] )
customer = customer.join(demo.select('cd_demo_sk', 'cd_gender', 'cd_marital_status', 'cd_credit_rating', 'cd_education_status', 'cd_dep_count'),
                                customer['c_current_hdemo_sk'] == demo['cd_demo_sk'] )
customer = customer.rename('c_customer_sk', 'customer_sk')
final_df = total_sales.join(customer, on='customer_sk')
session.use_database('tpcds_xgboost')
session.use_schema('demo')
final_df.write.mode('overwrite').save_as_table('feature_store')

Ora siamo pronti per addestrare il nostro modello usando una stored procedure. Snowflake ti mette semplicemente a disposizione tutti i popolari framework Python ML da utilizzare per la formazione ML tramite la nostra partnership con Anaconda . Non c'è bisogno di imparare una nuova libreria o sintassi.

from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, OneHotEncoder, MinMaxScaler
from sklearn.metrics import mean_squared_error
from sklearn.compose import ColumnTransformer
from xgboost import XGBRegressor
import joblib
import os

def train_model(session: snowflake.snowpark.Session) -> float:
    snowdf = session.table("feature_store")
    snowdf = snowdf.drop(['CUSTOMER_SK', 'C_CURRENT_HDEMO_SK', 'C_CURRENT_ADDR_SK', 'C_CUSTOMER_ID', 'CA_ADDRESS_SK', 'CD_DEMO_SK'])
    snowdf_train, snowdf_test = snowdf.random_split([0.8, 0.2], seed=82) 

    # save the train and test sets as time stamped tables in Snowflake 
    snowdf_train.write.mode("overwrite").save_as_table("tpcds_xgboost.demo.tpc_TRAIN")
    snowdf_test.write.mode("overwrite").save_as_table("tpcds_xgboost.demo.tpc_TEST")
    train_x = snowdf_train.drop("TOTAL_SALES").to_pandas() # drop labels for training set
    train_y = snowdf_train.select("TOTAL_SALES").to_pandas()
    test_x = snowdf_test.drop("TOTAL_SALES").to_pandas()
    test_y = snowdf_test.select("TOTAL_SALES").to_pandas()
    cat_cols = ['CA_ZIP', 'CD_GENDER', 'CD_MARITAL_STATUS', 'CD_CREDIT_RATING', 'CD_EDUCATION_STATUS']
    num_cols = ['C_BIRTH_YEAR', 'CD_DEP_COUNT']

    num_pipeline = Pipeline([
            ('imputer', SimpleImputer(strategy="median")),
            ('std_scaler', StandardScaler()),
        ])

    preprocessor = ColumnTransformer(
    transformers=[('num', num_pipeline, num_cols),
                  ('encoder', OneHotEncoder(handle_unknown="ignore"), cat_cols) ])

    pipe = Pipeline([('preprocessor', preprocessor), 
                        ('xgboost', XGBRegressor())])
    pipe.fit(train_x, train_y)

    test_preds = pipe.predict(test_x)
    rmse = mean_squared_error(test_y, test_preds)
    model_file = os.path.join('/tmp', 'model.joblib')
    joblib.dump(pipe, model_file)
    session.file.put(model_file, "@ml_models",overwrite=True)
    return rmse

train_model_sp = F.sproc(train_model, session=session, replace=True)
# Switch to Snowpark Optimized Warehouse for training and to run the stored proc
session.use_warehouse('snowpark_opt_wh')
train_model_sp(session=session)

import sys
import pandas as pd
import cachetools
import joblib
from snowflake.snowpark import types as T

session.add_import("@ml_models/model.joblib")  

features = [ 'C_BIRTH_YEAR', 'CA_ZIP', 'CD_GENDER', 'CD_MARITAL_STATUS', 'CD_CREDIT_RATING', 'CD_EDUCATION_STATUS', 'CD_DEP_COUNT']

@cachetools.cached(cache={})
def read_file(filename):
       import_dir = sys._xoptions.get("snowflake_import_directory")
       if import_dir:
              with open(os.path.join(import_dir, filename), 'rb') as file:
                     m = joblib.load(file)
                     return m

@F.pandas_udf(session=session, max_batch_size=10000, is_permanent=True, 
              stage_location='@ml_models', name="clv_xgboost_udf")
def predict(df:  T.PandasDataFrame[int, str, str, str, str, str, int]) -> T.PandasSeries[float]:
       m = read_file('model.joblib')       
       df.columns = features
       return m.predict(df)

inference_df = session.table('feature_store')
inference_df = inference_df.drop(['CUSTOMER_SK', 'C_CURRENT_HDEMO_SK', 'C_CURRENT_ADDR_SK', 'C_CUSTOMER_ID', 'CA_ADDRESS_SK', 'CD_DEMO_SK'])
inputs = inference_df.drop("TOTAL_SALES")
snowdf_results = inference_df.select(*inputs,
                    predict(*inputs).alias('PREDICTION'), 
                    (F.col('TOTAL_SALES')).alias('ACTUAL_SALES')
                    )
snowdf_results.write.mode('overwrite').save_as_table('predictions')

SELECT "C_BIRTH_YEAR", 
        "CA_ZIP", 
        "CD_GENDER", 
        "CD_MARITAL_STATUS", 
        "CD_CREDIT_RATING",
        "CD_EDUCATION_STATUS", 
        "CD_DEP_COUNT", 
 clv_xgboost_udf("C_BIRTH_YEAR", "CA_ZIP", "CD_GENDER", "CD_MARITAL_STATUS", "CD_CREDIT_RATING", "CD_EDUCATION_STATUS", "CD_DEP_COUNT") AS "PREDICTION", 
"TOTAL_SALES" AS "ACTUAL_SALES" 
FROM tpcds_xgboost.demo.feature_store

Come dico a tutti i miei clienti con cui lavoro, per favore non credetemi sulla parola. Provalo tu stesso, tutto il codice è disponibile per te qui . Ancora meglio di TPC-DS, provalo con alcuni dati e pipeline della tua organizzazione.