ML em Snowflake em escala com Snowpark Python e XGBoost

Ao longo do último ano e meio, nossa equipe trabalhou com centenas de clientes construindo no Snowflake com Snowpark e Python (aliás - confira esta postagem do meu colega Caleb Baechtold que compila todos os aprendizados sobre a operacionalização do Snowpark Python na produção) e um A pergunta que recebo com frequência é: “mas temos realmente big data, como isso funciona em escala?”.

Muitos dos exemplos/demonstrações que você encontrará on-line são ótimos para mostrar exemplos planejados simples que mostram a funcionalidade principal, mas poucos, se houver, mostram exemplos de como ele funciona no contexto de problemas de grande escala empresarial.

TPC-DS , é um conjunto de dados útil de grande escala que geralmente representa um esquema genérico de dados de negócios (não estou aqui para defender que TPC-DS, ou qualquer conjunto de dados seja o fim, todos sejam benchmarks de qualquer plataforma/ferramenta, eu apenas acho que pode ser um conjunto de dados útil como ponto de partida para testar coisas em escala). O Snowflake disponibiliza o TPC-DS em todas as contas do Snowflake como um compartilhamento de dados, tanto na edição de 10 TB quanto na de 100 TB (aliás, porque é exposto como um compartilhamento de dados, você como usuário não paga nenhum dos custos de armazenamento . Apenas para qualquer cálculo potencial que você usar para consultá-lo). A edição de 100 TB tem mais de 560 bilhões de linhas nas tabelas de fatos. A edição de 10 TB, 56 bilhões.

Podemos usar o Snowpark Python para criar uma solução de ML relativamente simples para um problema comercial comum de uma empresa hipotética como a TPC, que é “Quero ser capaz de prever o valor vitalício de meus clientes em todos os canais de vendas”. Usaremos a API Snowpark Python Dataframe para fazer a preparação de dados/engenharia de recursos, procedimentos armazenados e Snowpark Optimized Warehouses para treinamento e UDFs em lote para inferência, tudo sem que os dados precisem sair do Snowflake, utilizando os recursos de computação e a capacidade de escala do Snowflake.

Vamos começar com nosso código de preparação de dados/engenharia de recursos. De forma bastante simples, agregaremos as vendas por cliente em todos os canais . Em seguida, juntaremos isso às nossas tabelas de dimensão do cliente para obter possíveis recursos de 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')

Agora estamos prontos para treinar nosso modelo usando um procedimento armazenado. O Snowflake simplesmente disponibiliza todas as estruturas Python ML populares para você utilizar no treinamento de ML por meio de nossa parceria com a Anaconda . Não há necessidade de aprender uma nova biblioteca ou sintaxe.

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

Como digo a todos os meus clientes com quem trabalho, por favor, não acreditem na minha palavra. Tente você mesmo, todo o código está disponível para você aqui . Ainda melhor do que o TPC-DS, faça um test drive com alguns dos dados e pipelines de sua organização.