import lightgbm.dask as lgb_dask
from lightgbm import DaskLGBMClassifier
import dask_databricks
import dask.dataframe as dd

client = dask_databricks.get_client()

train_ddf = dd.read_parquet(train_filepath, storage_options=storage_options).repartition(npartitions=320)
eval_ddf = dd.read_parquet(eval_filepath, storage_options=storage_options).repartition(npartitions=320)
#have played with lots of different partitioning strategies

#cast categorical features as categorical types, categorize in the training data, apply that categorization to the eval data
train_ddf[categorical_features] = train_dff[categorical_features].astype('category').categorize()
category_mappings = {col: train_ddf[col].cat.categories for col in categorical_columns}#hold onto this guy for making predictions on unseen data
eval_ddf[categorical_features] = eval_dff[categorical_features].astype('category')
for col in categorical_columns:
  eval_ddf[col] = eval_ddf[col].cat.set_categories(category_mappings[col])

clf = DaskLGBMClassifier(
    client=client,
    objective="binary",
    max_depth=-1,
    num_leaves=5000,
    metric="binary_logloss",
    boosting_type="gbdt"
)

clf.fit(
    X=train_ddf[features], 
    y=train_ddf['target'], 
    eval_set=[(eval_ddf[features], eval_ddf['target'])],
    eval_names=['eval'],

)